Add support for queuing source blocks

Instead of sending execute requests immediately when executing
multiple source blocks in one go, they are queued such that when an
execute_request has received an execute_reply, the next one queued up
is sent.  If one of the source blocks errors during its execution, all
of the queued blocks get aborted.  That is they do not get sent to the
kernel.

closes #498

* README.org: Describe feature.

* test/jupyter-test.el (org-babel-src-block-queuing): New test.

* jupyter-monads.el (jupyter-message-publisher): Publish non-Jupyter
messages as is without processing them.
(jupyter-request): Only handle Jupyter messages in the relevant
subscriber.

* jupyter-org-client.el (jupyter-org-queue-requests): New
customization option to determine if client-side queuing should take
place.
(most-recent-request, last-queued-request): New slots to a
`jupyter-org-client`,
(jupyter-org-abort): New function that aborts a queued execute request.
(jupyter-org-maybe-queued): New function that does the queuing of
requests if it's necessary.

* ob-jupyter.el (org-babel-jupyter--execute): Use
`jupyter-org-maybe-queued`.

Author: nnicandro

README.org

@@ -581,6 +581,18 @@ This behavior can be suppressed by setting =jupyter-org-auto-connect=
 to =nil=.  In this case, a connection is attempted upon executing a
 source block, for example.
 
+*** Enable client-side queuing of requests
+
+If the customizable variable =jupyter-org-queue-requests= is non-nil,
+then perform client side queuing of source block execute requests.
+This means that when multiple requests are made, for example by
+executing a subtree, the requests are queued locally in Emacs instead
+of sending all the requests immediately to the kernel as would happen
+when =:async yes= is specified on all the source blocks.  It is only
+when one request finishes that the next is sent.  In addition, if any
+request fails all the queued requests that are meant to come after it
+are aborted and do not get sent to the kernel.
+
 ** Kernel/notebook server
 *** Managing live kernels
 

jupyter-monads.el

@@ -400,6 +400,8 @@ the callbacks."
     (jupyter-publisher
       (lambda (msg)
         (pcase (jupyter-message-type msg)
+          ;; Send what doesn't appear to be a message as is.
+          ((pred null) (jupyter-content msg))
           ;; A status message after a request goes idle means there is
           ;; a new request and there will, theoretically, be no more
           ;; messages for the idle one.
@@ -466,8 +468,10 @@ list, represents."
             (jupyter-subscribe
               (jupyter-subscriber
                 (lambda (msg)
-                  (let ((channel (plist-get msg :channel)))
-                    (jupyter-handle-message client channel msg)))))))
+                  ;; Only handle what looks to be a Jupyter message.
+                  (when (jupyter-message-type msg)
+                    (let ((channel (plist-get msg :channel)))
+                      (jupyter-handle-message client channel msg))))))))
         (cons req client)))))
 
 (provide 'jupyter-monads)

jupyter-org-client.el

@@ -72,6 +72,26 @@ session by, for example, executing a src-block."
   :group 'ob-jupyter
   :type 'boolean)
 
+(defcustom jupyter-org-queue-requests nil
+  "Whether or not source block evaluations should be queued.
+When this variable is nil and, for example, multiple source
+blocks are executed in rapid succession the underlying
+\"execute_request\" messages are sent to the kernel immediately
+and are queued on the kernel side so that when one of the source
+blocks raises an error, the kernel will typically just execute
+the next \"execute_request\" message queued up so the effect is
+that source blocks that come after the failed one are executed.
+Some may find this behavior undesirable.
+
+Instead, when this variable is non-nil, the \"execute_request\"
+messages of the source blocks are queued on the client side and
+whenever one of the source blocks raises an error, all of the
+queued \"execute_request\" messages are aborted and don't get
+sent to the kernel so the effect is that source blocks that come
+after the failed one are not executed."
+  :group 'ob-jupyter
+  :type 'boolean)
+
 (defcustom jupyter-org-resource-directory "./.ob-jupyter/"
   "Directory used to store automatically generated image files.
 See `jupyter-org-image-file-name'."
@@ -111,7 +131,16 @@ See also the docstring of `org-image-actual-width' for more details."
   "MIME types handled by Jupyter Org.")
 
 (defclass jupyter-org-client (jupyter-repl-client)
-  ())
+  ((most-recent-request
+    :type (or jupyter-request null)
+    :initform nil
+    :initarg :most-recent-request
+    :documentation "The most recently sent request.")
+   (last-queued-request
+    :type (or jupyter-request null)
+    :initform nil
+    :initarg :last-queued-request
+    :documentation "The last queued request.")))
 
 (cl-defstruct (jupyter-org-request
                (:include jupyter-request)
@@ -400,6 +429,96 @@ to."
       (org-with-point-at (jupyter-org-request-marker req)
         (run-hooks 'org-babel-after-execute-hook)))))
 
+;;; Queueing requests
+
+(defun jupyter-org-abort (req)
+  "Abort REQ.
+Set the request as being idle.  Remove any indication that REQ is
+a running execute_request from the Org buffer.  Publish an abort
+message down the chain of subscribers to the REQ's message
+publisher to indicate that any subsequent, queued, requests
+should also be aborted."
+  (setf (jupyter-request-idle-p req) t)
+  (let ((client (jupyter-request-client req)))
+    (when (eq (oref client last-queued-request) req)
+      (oset client last-queued-request nil)))
+  (jupyter-org--remove-overlay req)
+  (jupyter-org--clear-async-indicator req)
+  (let ((marker (jupyter-org-request-marker req)))
+    (message (format "Source block execution in %s at position %s canceled"
+                     (buffer-name (marker-buffer marker))
+                     (marker-position marker))))
+  (with-demoted-errors "Error while aborting subscribers: %S"
+    (jupyter-run-with-io
+        (jupyter-request-message-publisher req)
+      ;; Propagate the abort down the chain of queued requests.
+      (jupyter-publish 'abort)))
+  (jupyter-unsubscribe))
+
+(defun jupyter-org-maybe-queued (dreq)
+  "Return a monadic value that either sends or continues to delay DREQ.
+DREQ is an already delayed request, as returned by
+`jupyter-request' and friends.  When the value is bound to a
+client, using e.g. `jupyter-run-with-client', send DREQ if there
+are no queued requests otherwise queue DREQ.  The value returns
+the unboxed request contained in DREQ.
+
+If the variable `jupyter-org-queue-requests' is nil, just send
+the request immediately instead of attempting to queue it."
+  (if (not jupyter-org-queue-requests)
+      (jupyter-sent dreq)
+    (jupyter-mlet* ((client (jupyter-get-state))
+                    (req dreq))
+      (let* ((send
+              (lambda (req)
+                (jupyter-run-with-client client
+                  (jupyter-mlet* ((req (jupyter-sent
+                                        (jupyter-return req))))
+                    (oset client most-recent-request req)
+                    (jupyter-run-with-io
+                        (jupyter-request-message-publisher req)
+                      (jupyter-subscribe
+                        (jupyter-subscriber
+                          (lambda (msg)
+                            (when (or (eq msg 'abort)
+                                      (equal (jupyter-message-type msg) "execute_reply"))
+                              (when (eq (oref client most-recent-request) req)
+                                (oset client most-recent-request nil))
+                              (jupyter-unsubscribe))))))
+                    (when (eq (oref client last-queued-request) req)
+                      (oset client last-queued-request nil))
+                    (jupyter-return req)))))
+             (queue
+              ;; Subscribe REQ to the message publisher of QREQ such that
+              ;; REQ is sent or aborted when QREQ receives an
+              ;; execute_reply.
+              (lambda (qreq req)
+                (let ((pub (jupyter-request-message-publisher qreq)))
+                  (jupyter-run-with-io pub
+                    (jupyter-subscribe
+                      (jupyter-subscriber
+                        (lambda (msg)
+                          (if (eq msg 'abort)
+                              (jupyter-org-abort req)
+                            (pcase (jupyter-message-type msg)
+                              ("execute_reply"
+                               (jupyter-with-message-content msg (status)
+                                 (if (equal status "ok")
+                                     (funcall send req)
+                                   (jupyter-org-abort req)))
+                               (jupyter-unsubscribe))))))))))))
+        (let ((mreq (oref client most-recent-request)) qreq)
+          (cond
+           ((null mreq)
+            (funcall send req))
+           ((setq qreq (oref client last-queued-request))
+            (funcall queue qreq req)
+            (oset client last-queued-request req))
+           (t
+            (funcall queue mreq req)
+            (oset client last-queued-request req)))
+          (jupyter-return req))))))
+
 ;;; Completion in code blocks
 
 (defvar jupyter-org--src-block-cache nil

ob-jupyter.el

@@ -493,33 +493,34 @@ These parameters are handled internally."
 
 (defun org-babel-jupyter--execute (code async-p)
   (jupyter-run-with-client jupyter-current-client
-    (jupyter-mlet* ((req (jupyter-sent (jupyter-execute-request :code code))))
-      (jupyter-return
-        `(,req
-          ,(cond
-            (async-p
-             (when (bound-and-true-p org-export-current-backend)
-               (jupyter-add-idle-sync-hook
-                'org-babel-after-execute-hook req 'append))
-             (if (jupyter-org-request-inline-block-p req)
-                 org-babel-jupyter-async-inline-results-pending-indicator
-               ;; This returns the message ID of REQ as an indicator
-               ;; for the pending results.
-               (jupyter-org-pending-async-results req)))
-            (t
-             (jupyter-idle-sync req)
-             (if (jupyter-org-request-inline-block-p req)
-                 ;; When evaluating a source block synchronously, only the
-                 ;; :execute-result will be in `jupyter-org-request-results' since
-                 ;; stream results and any displayed data will be placed in a separate
-                 ;; buffer.
-                 (let ((el (jupyter-org-result
-                            req (car (jupyter-org-request-results req)))))
-                   (if (stringp el) el
-                     (org-element-property :value el)))
-               ;; This returns an Org formatted string of the collected
-               ;; results.
-               (jupyter-org-sync-results req)))))))))
+    (let ((dreq (jupyter-execute-request :code code)))
+      (jupyter-mlet* ((req (jupyter-org-maybe-queued dreq)))
+        (jupyter-return
+          `(,req
+            ,(cond
+              (async-p
+               (when (bound-and-true-p org-export-current-backend)
+                 (jupyter-add-idle-sync-hook
+                  'org-babel-after-execute-hook req 'append))
+               (if (jupyter-org-request-inline-block-p req)
+                   org-babel-jupyter-async-inline-results-pending-indicator
+                 ;; This returns the message ID of REQ as an indicator
+                 ;; for the pending results.
+                 (jupyter-org-pending-async-results req)))
+              (t
+               (jupyter-idle-sync req)
+               (if (jupyter-org-request-inline-block-p req)
+                   ;; When evaluating a source block synchronously, only the
+                   ;; :execute-result will be in `jupyter-org-request-results' since
+                   ;; stream results and any displayed data will be placed in a separate
+                   ;; buffer.
+                   (let ((el (jupyter-org-result
+                              req (car (jupyter-org-request-results req)))))
+                     (if (stringp el) el
+                       (org-element-property :value el)))
+                 ;; This returns an Org formatted string of the collected
+                 ;; results.
+                 (jupyter-org-sync-results req))))))))))
 
 (defvar org-babel-jupyter-current-src-block-params nil
   "The block parameters of the most recently executed Jupyter source block.")

test/jupyter-test.el

@@ -3048,6 +3048,60 @@ print(2)"
        (forward-line)
        (should (looking-at-p ": 6"))))))
 
+(ert-deftest org-babel-src-block-queuing  ()
+  :tags '(org)
+  (let ((jupyter-org-queue-requests t)
+        (src (apply #'format "\
+* Tree
+
+#+BEGIN_SRC jupyter-python :session %s :async yes
+1
+#+END_SRC
+
+#+BEGIN_SRC jupyter-python :session %s :async yes
+raise Exception(\"This is an error\")
+#+END_SRC
+
+#+BEGIN_SRC jupyter-python :session %s :async yes
+3
+#+END_SRC
+
+#+BEGIN_SRC jupyter-python :session %s :async yes
+4
+#+END_SRC
+
+#+BEGIN_SRC jupyter-python :session %s :async yes
+5
+#+END_SRC
+
+#+BEGIN_SRC jupyter-python :session %s :async yes
+6
+#+END_SRC
+"
+                    (cl-loop repeat 6 collect jupyter-org-test-session))))
+    (jupyter-org-test
+     (insert src)
+     (let ((check-result
+            (lambda (result)
+              (let ((pos (org-babel-where-is-src-block-result)))
+                (save-excursion
+                  (should pos)
+                  (goto-char pos)
+                  (forward-line)
+                  (should (equal (buffer-substring
+                                  (line-beginning-position)
+                                  (line-end-position))
+                                 result)))))))
+       (org-babel-execute-subtree)
+       (while (jupyter-org-request-at-point)
+         (sleep-for 0.1))
+       (funcall check-result "")
+       (org-previous-block 1)
+       (org-babel-execute-src-block)
+       (while (jupyter-org-request-at-point)
+         (sleep-for 0.1))
+       (funcall check-result ": 6")))))
+
 ;; Local Variables:
 ;; byte-compile-warnings: (unresolved obsolete lexical)
 ;; eval: (and (functionp 'aggressive-indent-mode) (aggressive-indent-mode -1))