merge @SpencerPark commit 710233f

From SpencerPark <[email protected]>:
    Lock writes to the sockets to fix a bug with interleaved publishes.

Author: cosmos72

kernel.go

@@ -36,14 +36,20 @@ type ConnectionInfo struct {
 	IP              string `json:"ip"`
 }
 
+// Socket wraps a zmq socket with a lock which should be used to control write access.
+type Socket struct {
+	Socket *zmq.Socket
+	Lock   *sync.Mutex
+}
+
 // SocketGroup holds the sockets needed to communicate with the kernel,
 // and the key for message signing.
 type SocketGroup struct {
-	ShellSocket   *zmq.Socket
-	ControlSocket *zmq.Socket
-	StdinSocket   *zmq.Socket
-	IOPubSocket   *zmq.Socket
-	HBSocket      *zmq.Socket
+	ShellSocket   Socket
+	ControlSocket Socket
+	StdinSocket   Socket
+	IOPubSocket   Socket
+	HBSocket      Socket
 	Key           []byte
 }
 
@@ -85,6 +91,13 @@ const (
 	kernelIdle     = "idle"
 )
 
+// RunWithSocket invokes the `run` function after acquiring the `Socket.Lock` and releases the lock when done.
+func (s *Socket) RunWithSocket(run func(socket *zmq.Socket) error) error {
+	s.Lock.Lock()
+	defer s.Lock.Unlock()
+	return run(s.Socket)
+}
+
 // runKernel is the main entry point to start the kernel.
 func runKernel(connectionFile string) {
 
@@ -127,9 +140,9 @@ func runKernel(connectionFile string) {
 	// TODO gracefully shutdown the heartbeat handler on kernel shutdown by closing the chan returned by startHeartbeat.
 
 	poller := zmq.NewPoller()
-	poller.Add(sockets.ShellSocket, zmq.POLLIN)
-	poller.Add(sockets.StdinSocket, zmq.POLLIN)
-	poller.Add(sockets.ControlSocket, zmq.POLLIN)
+	poller.Add(sockets.ShellSocket.Socket, zmq.POLLIN)
+	poller.Add(sockets.StdinSocket.Socket, zmq.POLLIN)
+	poller.Add(sockets.ControlSocket.Socket, zmq.POLLIN)
 
 	// msgParts will store a received multipart message.
 	var msgParts [][]byte
@@ -147,8 +160,8 @@ func runKernel(connectionFile string) {
 			switch socket := item.Socket; socket {
 
 			// Handle shell messages.
-			case sockets.ShellSocket:
-				msgParts, err = sockets.ShellSocket.RecvMessageBytes(0)
+			case sockets.ShellSocket.Socket:
+				msgParts, err = sockets.ShellSocket.Socket.RecvMessageBytes(0)
 				if err != nil {
 					log.Println(err)
 				}
@@ -162,12 +175,12 @@ func runKernel(connectionFile string) {
 				handleShellMsg(ir, msgReceipt{msg, ids, sockets})
 
 				// TODO Handle stdin socket.
-			case sockets.StdinSocket:
-				sockets.StdinSocket.RecvMessageBytes(0)
+			case sockets.StdinSocket.Socket:
+				sockets.StdinSocket.Socket.RecvMessageBytes(0)
 
 				// Handle control messages.
-			case sockets.ControlSocket:
-				msgParts, err = sockets.ControlSocket.RecvMessageBytes(0)
+			case sockets.ControlSocket.Socket:
+				msgParts, err = sockets.ControlSocket.Socket.RecvMessageBytes(0)
 				if err != nil {
 					log.Println(err)
 					return
@@ -200,46 +213,51 @@ func prepareSockets(connInfo ConnectionInfo) (SocketGroup, error) {
 
 	// Create the shell socket, a request-reply socket that may receive messages from multiple frontend for
 	// code execution, introspection, auto-completion, etc.
-	sg.ShellSocket, err = context.NewSocket(zmq.ROUTER)
+	sg.ShellSocket.Socket, err = context.NewSocket(zmq.ROUTER)
+	sg.ShellSocket.Lock = &sync.Mutex{}
 	if err != nil {
 		return sg, err
 	}
 
 	// Create the control socket. This socket is a duplicate of the shell socket where messages on this channel
 	// should jump ahead of queued messages on the shell socket.
-	sg.ControlSocket, err = context.NewSocket(zmq.ROUTER)
+	sg.ControlSocket.Socket, err = context.NewSocket(zmq.ROUTER)
+	sg.ControlSocket.Lock = &sync.Mutex{}
 	if err != nil {
 		return sg, err
 	}
 
 	// Create the stdin socket, a request-reply socket used to request user input from a front-end. This is analogous
 	// to a standard input stream.
-	sg.StdinSocket, err = context.NewSocket(zmq.ROUTER)
+	sg.StdinSocket.Socket, err = context.NewSocket(zmq.ROUTER)
+	sg.StdinSocket.Lock = &sync.Mutex{}
 	if err != nil {
 		return sg, err
 	}
 
 	// Create the iopub socket, a publisher for broadcasting data like stdout/stderr output, displaying execution
 	// results or errors, kernel status, etc. to connected subscribers.
-	sg.IOPubSocket, err = context.NewSocket(zmq.PUB)
+	sg.IOPubSocket.Socket, err = context.NewSocket(zmq.PUB)
+	sg.IOPubSocket.Lock = &sync.Mutex{}
 	if err != nil {
 		return sg, err
 	}
 
 	// Create the heartbeat socket, a request-reply socket that only allows alternating recv-send (request-reply)
 	// calls. It should echo the byte strings it receives to let the requester know the kernel is still alive.
-	sg.HBSocket, err = context.NewSocket(zmq.REP)
+	sg.HBSocket.Socket, err = context.NewSocket(zmq.REP)
+	sg.HBSocket.Lock = &sync.Mutex{}
 	if err != nil {
 		return sg, err
 	}
 
 	// Bind the sockets.
 	address := fmt.Sprintf("%v://%v:%%v", connInfo.Transport, connInfo.IP)
-	sg.ShellSocket.Bind(fmt.Sprintf(address, connInfo.ShellPort))
-	sg.ControlSocket.Bind(fmt.Sprintf(address, connInfo.ControlPort))
-	sg.StdinSocket.Bind(fmt.Sprintf(address, connInfo.StdinPort))
-	sg.IOPubSocket.Bind(fmt.Sprintf(address, connInfo.IOPubPort))
-	sg.HBSocket.Bind(fmt.Sprintf(address, connInfo.HBPort))
+	sg.ShellSocket.Socket.Bind(fmt.Sprintf(address, connInfo.ShellPort))
+	sg.ControlSocket.Socket.Bind(fmt.Sprintf(address, connInfo.ControlPort))
+	sg.StdinSocket.Socket.Bind(fmt.Sprintf(address, connInfo.StdinPort))
+	sg.IOPubSocket.Socket.Bind(fmt.Sprintf(address, connInfo.IOPubPort))
+	sg.HBSocket.Socket.Bind(fmt.Sprintf(address, connInfo.HBPort))
 
 	// Set the message signing key.
 	sg.Key = []byte(connInfo.Key)
@@ -496,7 +514,7 @@ func handleShutdownRequest(receipt msgReceipt) {
 // startHeartbeat starts a go-routine for handling heartbeat ping messages sent over the given `hbSocket`. The `wg`'s
 // `Done` method is invoked after the thread is completely shutdown. To request a shutdown the returned `shutdown` channel
 // can be closed.
-func startHeartbeat(hbSocket *zmq.Socket, wg *sync.WaitGroup) (shutdown chan struct{}) {
+func startHeartbeat(hbSocket Socket, wg *sync.WaitGroup) (shutdown chan struct{}) {
 	quit := make(chan struct{})
 
 	// Start the handler that will echo any received messages back to the sender.
@@ -506,7 +524,7 @@ func startHeartbeat(hbSocket *zmq.Socket, wg *sync.WaitGroup) (shutdown chan str
 
 		// Create a `Poller` to check for incoming messages.
 		poller := zmq.NewPoller()
-		poller.Add(hbSocket, zmq.POLLIN)
+		poller.Add(hbSocket.Socket, zmq.POLLIN)
 
 		for {
 			select {
@@ -521,16 +539,22 @@ func startHeartbeat(hbSocket *zmq.Socket, wg *sync.WaitGroup) (shutdown chan str
 
 				// If there is at least 1 message waiting then echo it.
 				if len(pingEvents) > 0 {
-					// Read a message from the heartbeat channel as a simple byte string.
-					pingMsg, err := hbSocket.RecvBytes(0)
-					if err != nil {
-						log.Fatalf("Error reading heartbeat ping bytes: %v\n", err)
-					}
-
-					// Send the received byte string back to let the front-end know that the kernel is alive.
-					if _, err = hbSocket.SendBytes(pingMsg, 0); err != nil {
-						log.Printf("Error sending heartbeat pong bytes: %b\n", err)
-					}
+					hbSocket.RunWithSocket(func(echo *zmq.Socket) error {
+						// Read a message from the heartbeat channel as a simple byte string.
+						pingMsg, err := echo.RecvBytes(0)
+						if err != nil {
+							log.Fatalf("Error reading heartbeat ping bytes: %v\n", err)
+							return err
+						}
+
+						// Send the received byte string back to let the front-end know that the kernel is alive.
+						if _, err = echo.SendBytes(pingMsg, 0); err != nil {
+							log.Printf("Error sending heartbeat pong bytes: %b\n", err)
+							return err
+						}
+
+						return nil
+					})
 				}
 			}
 		}

messages.go

@@ -189,7 +189,9 @@ func (receipt *msgReceipt) Publish(msgType string, content interface{}) error {
 	}
 
 	msg.Content = content
-	return receipt.SendResponse(receipt.Sockets.IOPubSocket, msg)
+	return receipt.Sockets.IOPubSocket.RunWithSocket(func(iopub *zmq.Socket) error {
+		return receipt.SendResponse(iopub, msg)
+	})
 }
 
 // Reply creates a new ComposedMsg and sends it back to the return identities over the
@@ -202,7 +204,9 @@ func (receipt *msgReceipt) Reply(msgType string, content interface{}) error {
 	}
 
 	msg.Content = content
-	return receipt.SendResponse(receipt.Sockets.ShellSocket, msg)
+	return receipt.Sockets.ShellSocket.RunWithSocket(func(shell *zmq.Socket) error {
+		return receipt.SendResponse(shell, msg)
+	})
 }
 
 // newTextMIMEDataBundle creates a bundledMIMEData that only contains a text representation described