Fix deadlock when progressive invocation handler errors (gammazero#319)

Author: Skinner927

client/client.go

@@ -69,7 +69,6 @@ type Client struct {
 	closed bool
 
 	routerGoodbye *wamp.Goodbye
-	idGen         *wamp.SyncIDGen
 }
 
 // InvokeResult represents the result of invoking a procedure.
@@ -256,7 +255,6 @@ func NewClient(p wamp.Peer, cfg Config) (*Client, error) {
 		log:        cfg.Logger,
 		debug:      cfg.Debug,
 		cancelMode: wamp.CancelModeKillNoWait,
-		idGen:      new(wamp.SyncIDGen),
 	}
 	c.ctx, c.cancel = context.WithCancel(context.Background())
 	go c.run() // start the core goroutine
@@ -313,7 +311,7 @@ func (c *Client) Subscribe(topic string, fn EventHandler, options wamp.Dict) err
 	if options == nil {
 		options = wamp.Dict{}
 	}
-	id := c.idGen.Next()
+	id := c.sess.IdGen.Next()
 	c.expectReply(id)
 	c.sess.Send() <- &wamp.Subscribe{
 		Request: id,
@@ -384,7 +382,7 @@ func (c *Client) Unsubscribe(topic string) error {
 		return ErrNotConn
 	}
 
-	id := c.idGen.Next()
+	id := c.sess.IdGen.Next()
 	c.expectReply(id)
 	c.sess.Send() <- &wamp.Unsubscribe{
 		Request:      id,
@@ -453,7 +451,7 @@ func (c *Client) Publish(topic string, options wamp.Dict, args wamp.List, kwargs
 		return ErrNotConn
 	}
 
-	id := c.idGen.Next()
+	id := c.sess.IdGen.Next()
 
 	var pubAck bool
 	if options == nil {
@@ -570,7 +568,7 @@ func (c *Client) Register(procedure string, fn InvocationHandler, options wamp.D
 	if !c.Connected() {
 		return ErrNotConn
 	}
-	id := c.idGen.Next()
+	id := c.sess.IdGen.Next()
 	c.expectReply(id)
 	if options == nil {
 		options = wamp.Dict{}
@@ -636,7 +634,7 @@ func (c *Client) Unregister(procedure string) error {
 		return ErrNotConn
 	}
 
-	id := c.idGen.Next()
+	id := c.sess.IdGen.Next()
 	c.expectReply(id)
 	c.sess.Send() <- &wamp.Unregister{
 		Request:      id,
@@ -758,7 +756,7 @@ func (c *Client) Call(ctx context.Context, procedure string, options wamp.Dict,
 		}()
 	}
 
-	id := c.idGen.Next()
+	id := c.sess.IdGen.Next()
 	c.expectReply(id)
 	message := &wamp.Call{
 		Request:   id,
@@ -857,7 +855,7 @@ func (c *Client) CallProgressive(ctx context.Context, procedure string, sendProg
 		}()
 	}
 
-	id := c.idGen.Next()
+	id := c.sess.IdGen.Next()
 	c.expectReply(id)
 	message := &wamp.Call{
 		Request:   id,
@@ -1445,14 +1443,17 @@ func (c *Client) runReceiveFromRouter(msg wamp.Message) bool {
 	if c.debug {
 		c.log.Println("Client", c.sess, "received", msg.MessageType())
 	}
+
 	switch msg := msg.(type) {
 	case *wamp.Event:
 		c.runHandleEvent(msg)
 
 	case *wamp.Invocation:
 		c.runHandleInvocation(msg)
+		c.sess.UpdateLastRecvID(msg.Request)
 	case *wamp.Interrupt:
 		c.runHandleInterrupt(msg)
+		c.sess.UpdateLastRecvID(msg.Request)
 
 	case *wamp.Registered:
 		c.runSignalReply(msg, msg.Request)
@@ -1535,6 +1536,38 @@ func (c *Client) runHandleEvent(msg *wamp.Event) {
 	handler(msg)
 }
 
+func (c *Client) cleanupInvHandlersQueue(cliInvocation clientInvocation) {
+	c.sess.Lock()
+	defer c.sess.Unlock()
+
+	handlerQueue, _ := c.invHandlersQueues[cliInvocation]
+	delete(c.invHandlersQueues, cliInvocation)
+	delete(c.invHandlersCtxs, cliInvocation)
+
+	if nil == handlerQueue {
+		return
+	}
+	if c.debug {
+		c.log.Println("Running cleanupInvHandlersQueue cleanup for", cliInvocation)
+	}
+
+	// Try to get any remaining values off the chan (in case anyone is blocked)
+	for {
+		select {
+		case msg := <-handlerQueue:
+			if nil == msg {
+				return // chan closed
+			}
+			continue
+		default:
+			return
+		}
+	}
+	// Don't bother closing in the very rare event that someone still has
+	// a reference and tries to send
+	// close(c.invHandlersQueues[cliInvocation])
+}
+
 // runHandleInvocation processes an INVOCATION message from the router
 // requesting a call to a registered RPC procedure.
 func (c *Client) runHandleInvocation(msg *wamp.Invocation) {
@@ -1616,6 +1649,19 @@ func (c *Client) runHandleInvocation(msg *wamp.Invocation) {
 	handlerQueue, queueExists := c.invHandlersQueues[cliInvocation]
 	ctx := c.invHandlersCtxs[cliInvocation]
 	if !queueExists {
+		// Only create the queue if this is a new request
+		if !c.sess.UpdateLastRecvIDCallerHasSessionLock(reqID) {
+			c.sess.Unlock()
+			if c.debug {
+				c.log.Println("Ignoring Invocation with expired reqID=", reqID)
+			}
+			// discard silently
+			return
+		}
+
+		if c.debug {
+			c.log.Println("Creating new handlerQueue reqID=", reqID)
+		}
 		handlerQueue = make(chan *wamp.Invocation, 1)
 		c.invHandlersQueues[cliInvocation] = handlerQueue
 
@@ -1638,7 +1684,6 @@ func (c *Client) runHandleInvocation(msg *wamp.Invocation) {
 			ctx = context.WithValue(ctx, invocationIDCtxKey{}, reqID)
 		}
 	}
-
 	c.sess.Unlock()
 
 	handlerQueue <- msg
@@ -1652,21 +1697,36 @@ func (c *Client) runHandleInvocation(msg *wamp.Invocation) {
 			// Otherwise, canceling the call will leak the goroutine that is
 			// blocked forever waiting to send the result to the channel.
 			resChan := make(chan InvokeResult, 1)
+
+			// Start goroutine to process inbound Invocation messages
 			go func() {
-				for msg := range handlerQueue {
-
-					if isInProgress, _ := msg.Details[wamp.OptProgress].(bool); !isInProgress {
-						c.sess.Lock()
-						close(c.invHandlersQueues[cliInvocation])
-						delete(c.invHandlersQueues, cliInvocation)
-						delete(c.invHandlersCtxs, cliInvocation)
-						c.sess.Unlock()
-					}
+				defer c.cleanupInvHandlersQueue(cliInvocation)
+
+				processMessages := true
+				for processMessages {
+					select {
+					case msg := <-handlerQueue:
+						if msg == nil { // chan closed
+							return
+						}
+						if isInProgress, _ := msg.Details[wamp.OptProgress].(bool); !isInProgress {
+							c.cleanupInvHandlersQueue(cliInvocation)
+							processMessages = false
+						}
 
-					// The Context is passed into the handler to tell the client
-					// application to stop whatever it is doing if it cares to pay
-					// attention.
-					resChan <- handler(ctx, msg)
+						// The Context is passed into the handler to tell the client
+						// application to stop whatever it is doing if it cares to pay
+						// attention.
+						result := handler(ctx, msg)
+						resChan <- result
+						if result.Err != "" && result.Err != wamp.InternalProgressiveOmitResult {
+							processMessages = false
+						}
+					case <-c.Done():
+						return
+					case <-ctx.Done():
+						return
+					}
 				}
 			}()
 

client/client_test.go

@@ -30,6 +30,9 @@ const (
 
 	testTopic  = "test.topic1"
 	testTopic2 = "test.topic2"
+
+	debugClientEnv = "TEST_DEBUG_CLIENT"
+	debugRouterEnv = "TEST_DEBUG_ROUTER"
 )
 
 var logger stdlog.StdLog
@@ -47,6 +50,7 @@ func checkGoLeaks(t *testing.T) {
 func getTestRouter(t *testing.T, realmConfig *router.RealmConfig) router.Router {
 	config := &router.Config{
 		RealmConfigs: []*router.RealmConfig{realmConfig},
+		Debug:        os.Getenv(debugRouterEnv) != "",
 	}
 	r, err := router.NewRouter(config, logger)
 	require.NoError(t, err)
@@ -95,7 +99,7 @@ func newTestClientConfig(realmName string, fns ...clientConfigMutator) *Config {
 		Realm:           realmName,
 		ResponseTimeout: 500 * time.Millisecond,
 		Logger:          logger,
-		Debug:           false,
+		Debug:           os.Getenv(debugClientEnv) != "",
 	}
 	for _, fn := range fns {
 		fn(clientConfig)
@@ -704,6 +708,97 @@ func TestProgressiveCallInvocations(t *testing.T) {
 	require.NoError(t, err)
 }
 
+// Tests that a callee can return error while caller IsInProgress
+// Test for #319 to ensure messages in-flight do not re-open the handlerQueue
+// or call the handler function after the callee has errored.
+func TestProgressiveCallInvocationCalleeError(t *testing.T) {
+	// Connect two clients to the same server
+	// t.Setenv(debugRouterEnv, "1")
+	t.Setenv(debugClientEnv, "1")
+	callee, caller, rooter := connectedTestClients(t)
+
+	const forcedError = wamp.URI("error.forced")
+	moreArgsSent := make(chan struct{})
+	errorRaised := false
+
+	invocationHandler := func(ctx context.Context, inv *wamp.Invocation) InvokeResult {
+		switch inv.Arguments[0].(int) {
+		case 1:
+			// Eat the first arg
+			t.Log("n=1 Returning OmitResult")
+			return InvokeResult{Err: wamp.InternalProgressiveOmitResult}
+		case 2:
+			t.Log("n=2 Waiting for moreArgsSent")
+			// Wait till the 4th arg is sent which means 3 should already
+			// be waiting
+			<-moreArgsSent
+			time.Sleep(100 * time.Millisecond)
+			errorRaised = true
+			t.Log("n=2 Returning error (as expected)")
+			// Error
+			return InvokeResult{Err: forcedError}
+		default:
+			// BUG: The handler function should never be called again
+			t.Error("Handler should not have been called after error returned")
+			return InvokeResult{Err: wamp.ErrInvalidArgument}
+		}
+	}
+
+	const procName = "nexus.test.progprocerr"
+
+	// Register procedure
+	err := callee.Register(procName, invocationHandler, nil)
+	require.NoError(t, err)
+
+	// Test calling the procedure.
+	callArgs := [...]int{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
+	ctx := context.Background()
+
+	sendCount := 0
+	sendProgDataCb := func(ctx context.Context) (options wamp.Dict, args wamp.List, kwargs wamp.Dict, err error) {
+		options = wamp.Dict{
+			wamp.OptProgress: sendCount < (len(callArgs) - 1),
+		}
+
+		args = wamp.List{callArgs[sendCount]}
+		sendCount++
+
+		// signal the handler should return its error
+		if 4 == sendCount {
+			close(moreArgsSent)
+		}
+		t.Logf("Sending n=%v", sendCount)
+
+		return options, args, nil, nil
+	}
+
+	result, err := caller.CallProgressive(ctx, procName, sendProgDataCb, nil)
+	require.Error(t, err, "Expected call to return an error")
+	require.Nil(t, result, "Expected call to return no result")
+	var rErr RPCError
+	if errors.As(err, &rErr) {
+		require.Equal(t, forcedError, rErr.Err.Error, "Unexpected error URI")
+	} else {
+		t.Error("Unexpected error type")
+	}
+	require.GreaterOrEqual(t, sendCount, 4)
+	require.True(t, errorRaised, "Error was never raised in handler")
+
+	// #319: Show deadlock
+	t.Log("Closing rooter")
+	rooter.Close()
+	t.Log("Closing caller")
+	require.NoError(t, caller.Close())
+	goleak.VerifyNone(t)
+	t.Log("Closing callee")
+	// #319: We never get past here
+	require.NoError(t, callee.Close())
+
+	t.Log("All closed")
+	goleak.VerifyNone(t)
+	t.Log("Done")
+}
+
 func TestProgressiveCallsAndResults(t *testing.T) {
 	// Connect two clients to the same server
 	callee, caller, _ := connectedTestClients(t)

wamp/idgen.go

@@ -6,15 +6,15 @@ import (
 	"time"
 )
 
-const maxID int64 = 1 << 53
+const MaxID uint64 = 1 << 53
 
 func init() {
 	rand.Seed(time.Now().UnixNano())
 }
 
 // NewID generates a random WAMP ID.
 func GlobalID() ID {
-	return ID(rand.Int63n(maxID)) //nolint:gosec
+	return ID(rand.Int63n(int64(MaxID))) //nolint:gosec
 }
 
 // IDGen is generator for WAMP request IDs.  Create with new(IDGen).
@@ -29,13 +29,13 @@ func GlobalID() ID {
 //
 // See https://github.com/wamp-proto/wamp-proto/blob/master/spec/basic.md#ids
 type IDGen struct {
-	next int64
+	next uint64
 }
 
 // Next returns next ID.
 func (g *IDGen) Next() ID {
 	g.next++
-	if g.next > maxID {
+	if g.next > MaxID {
 		g.next = 1
 	}
 	return ID(g.next)

wamp/idgen_test.go

@@ -26,7 +26,7 @@ func TestIDGen(t *testing.T) {
 	require.Equal(t, ID(2), id2, errMsg)
 	require.Equal(t, ID(3), id3, errMsg)
 
-	idgen.next = int64(1) << 53
+	idgen.next = uint64(1) << 53
 	id1 = idgen.Next()
 	require.Equal(t, ID(1), id1, "Sequential IDs should wrap at 1 << 53")
 }