From bf7eee0f658bf0b70e257777dcdebdf3625306e6 Mon Sep 17 00:00:00 2001
From: Matthew Flatt <mflatt@racket-lang.org>
Date: Wed, 14 Nov 2018 16:44:01 -0700
Subject: [PATCH] thread: fix atomic-timeout-callback implementation

Closes racket/drracket#238
---
 .../tests/racket/try-atomic.rkt               | 21 +++++++++++++++++++
 racket/src/thread/Makefile                    |  3 ++-
 racket/src/thread/schedule.rkt                | 10 ++++++++-
 racket/src/thread/thread.rkt                  | 13 ++++++++++--
 4 files changed, 43 insertions(+), 4 deletions(-)
 create mode 100644 pkgs/racket-test-core/tests/racket/try-atomic.rkt

diff --git a/pkgs/racket-test-core/tests/racket/try-atomic.rkt b/pkgs/racket-test-core/tests/racket/try-atomic.rkt
new file mode 100644
index 0000000000..a7d90acb75
--- /dev/null
+++ b/pkgs/racket-test-core/tests/racket/try-atomic.rkt
@@ -0,0 +1,21 @@
+#lang racket/base
+(require ffi/unsafe/try-atomic)
+
+(define ch (make-channel))
+(define done? #f)
+
+(define (check v expect)
+  (unless (equal? v expect)
+    (error 'check "failed: ~s vs. ~s" v expect)))
+
+(check
+ (call-as-nonatomic-retry-point
+  (lambda ()
+    (try-atomic
+     (lambda ()
+       (sync/timeout 0.25 ch)
+      (set! done? #t))
+     'no)))
+ 'no)
+
+(check done? #t)
diff --git a/racket/src/thread/Makefile b/racket/src/thread/Makefile
index a17293f927..b0297cb714 100644
--- a/racket/src/thread/Makefile
+++ b/racket/src/thread/Makefile
@@ -19,7 +19,8 @@ GLOBALS = --no-global \
           ++global-ok the-sandman \
           ++global-ok sync-on-channel \
           ++global-ok post-shutdown-action \
-          ++global-ok get-subprocesses-time
+          ++global-ok get-subprocesses-time \
+          ++global-ok force-atomic-timeout-callback
 
 GENERATE_ARGS = -t main.rkt --submod main \
                 --check-depends $(BUILDDIR)compiled/thread-dep.rktd \
diff --git a/racket/src/thread/schedule.rkt b/racket/src/thread/schedule.rkt
index 67892e4294..e2d8b84f43 100644
--- a/racket/src/thread/schedule.rkt
+++ b/racket/src/thread/schedule.rkt
@@ -82,7 +82,7 @@
           (check-for-break)
           (when atomic-timeout-callback
             (when (positive? (current-atomic))
-              (atomic-timeout-callback))))
+              (atomic-timeout-callback #f))))
         (lambda args
           (start-implicit-atomic-mode)
           (accum-cpu-time! t)
@@ -235,6 +235,14 @@
     atomic-timeout-callback
     (set! atomic-timeout-callback cb)))
 
+
+(void (set-force-atomic-timeout-callback!
+       (lambda ()
+         (and atomic-timeout-callback
+              (begin
+                (atomic-timeout-callback #t)
+                #t)))))
+
 ;; ----------------------------------------
 
 (define check-place-activity void)
diff --git a/racket/src/thread/thread.rkt b/racket/src/thread/thread.rkt
index ab32235b34..ad2e98f667 100644
--- a/racket/src/thread/thread.rkt
+++ b/racket/src/thread/thread.rkt
@@ -82,6 +82,8 @@
            check-for-break
            current-break-suspend
 
+           set-force-atomic-timeout-callback!
+
            break-max))
 
 ;; Exports needed by "place.rkt":
@@ -387,6 +389,10 @@
   (define sleeping (sandman-add-sleeping-thread! t ext-events))
   (set-thread-sleeping! t sleeping))
 
+(define force-atomic-timeout-callback void)
+(define (set-force-atomic-timeout-callback! proc)
+  (set! force-atomic-timeout-callback proc))
+
 ;; in atomic mode
 ;; Removes a thread from its thread group, so it won't be scheduled;
 ;; returns a thunk to be called in out of atomic mode to swap out the
@@ -405,8 +411,11 @@
   ;; by a custodian callback
   (lambda ()
     (when (eq? t (current-thread))
-      (when (positive? (current-atomic))
-        (internal-error "attempt to deschedule the current thread in atomic mode"))
+      (let loop ()
+        (when (positive? (current-atomic))
+          (if (force-atomic-timeout-callback)
+              (loop)
+              (internal-error "attempt to deschedule the current thread in atomic mode"))))
       (engine-block)
       (check-for-break))))