github.com/yunabe/lgo@v0.0.0-20190709125917-42c42d410fdf/jupyter/gojupyterscaffold/execute.go (about)

     1  // executeQueue executes execute_requests sequentially.
     2  // Although shellSocket in this package accepts request on a shell socket in parallel,
     3  // we shold not handle multiple execute_requests in parallel because the jupyter client sends
     4  // successive execute_requests to servers before previous execute_requests finishes.
     5  // A jupyter kernel is responsible to handle multiple execute_requests sequentially and
     6  // abort them if one of them fails.
     7  
     8  package gojupyterscaffold
     9  
    10  import (
    11  	"context"
    12  	"errors"
    13  	"fmt"
    14  )
    15  
    16  const executeQueueSize = 1 << 8
    17  
    18  type executeQueueItem struct {
    19  	req  *message
    20  	sock *shellSocket
    21  }
    22  
    23  type executeQueue struct {
    24  	serverCtx  context.Context
    25  	queue      chan *executeQueueItem
    26  	iopub      *iopubSocket
    27  	handlers   RequestHandlers
    28  	currentCtx *contextAndCancel
    29  }
    30  
    31  func newExecuteQueue(ctx context.Context, iopub *iopubSocket, handlers RequestHandlers) *executeQueue {
    32  	return &executeQueue{
    33  		serverCtx: ctx,
    34  		queue:     make(chan *executeQueueItem, executeQueueSize),
    35  		iopub:     iopub,
    36  		handlers:  handlers,
    37  	}
    38  }
    39  
    40  func (q *executeQueue) push(req *message, sock *shellSocket) {
    41  	q.queue <- &executeQueueItem{req, sock}
    42  }
    43  
    44  // abortQueue aborts requests in the queue.
    45  // c.f. _abort_queue in https://github.com/ipython/ipykernel/blob/master/ipykernel/kernelbase.py
    46  func (q *executeQueue) abortQueue() {
    47  loop:
    48  	for {
    49  		var item *executeQueueItem
    50  		select {
    51  		case item = <-q.queue:
    52  		default:
    53  			break loop
    54  		}
    55  		err := q.iopub.WithOngoingContext(func(ctx context.Context) error {
    56  			res := newMessageWithParent(item.req)
    57  			res.Header.MsgType = "execute_reply"
    58  			res.Content = &ExecuteResult{
    59  				Status: "abort",
    60  			}
    61  			if err := item.sock.pushResult(res); err != nil {
    62  				return fmt.Errorf("Failed to send execute_reply: %v", err)
    63  			}
    64  			return nil
    65  		}, item.req)
    66  		if err != nil {
    67  			logger.Errorf("Failed to abort a execute request: %v", err)
    68  		}
    69  	}
    70  }
    71  
    72  func (q *executeQueue) cancelCurrent() {
    73  	cur := q.currentCtx
    74  	if cur != nil {
    75  		cur.cancel()
    76  	}
    77  }
    78  
    79  // loop executes execute_requests sequentially.
    80  func (q *executeQueue) loop() {
    81  	var errStatusError = errors.New("execute status error")
    82  loop:
    83  	for {
    84  		var item *executeQueueItem
    85  		select {
    86  		case item = <-q.queue:
    87  		case <-q.serverCtx.Done():
    88  			break loop
    89  		}
    90  
    91  		exReq := item.req.Content.(*ExecuteRequest)
    92  		err := q.iopub.WithOngoingContext(func(ctx context.Context) error {
    93  			cur, cancel := context.WithCancel(ctx)
    94  			q.currentCtx = &contextAndCancel{cur, cancel}
    95  			defer func() {
    96  				cancel()
    97  				q.currentCtx = nil
    98  			}()
    99  			result := q.handlers.HandleExecuteRequest(
   100  				cur,
   101  				exReq,
   102  				func(name, text string) {
   103  					q.iopub.sendStream(name, text, item.req)
   104  				}, func(data *DisplayData, update bool) {
   105  					q.iopub.sendDisplayData(data, item.req, update)
   106  				})
   107  			res := newMessageWithParent(item.req)
   108  			res.Header.MsgType = "execute_reply"
   109  			res.Content = &result
   110  			if err := item.sock.pushResult(res); err != nil {
   111  				logger.Errorf("Failed to send execute_reply: %v", err)
   112  			}
   113  			if result.Status == "error" {
   114  				return errStatusError
   115  			}
   116  			return nil
   117  		}, item.req)
   118  		if err != nil {
   119  			if err != errStatusError {
   120  				logger.Errorf("Failed to handle a execute_request: %v", err)
   121  			}
   122  			if exReq.StopOnError {
   123  				q.abortQueue()
   124  			}
   125  		}
   126  	}
   127  }