jsonrpc: Handle closing channels

lib/jsonrpc/client.go

@@ -45,7 +45,8 @@ type clientRequest struct {
 	req   request
 	ready chan clientResponse
 
-	retCh func() func([]byte, bool)
+	// retCh provides a context and sink for handling incoming channel messages
+	retCh func() (context.Context, func([]byte, bool))
 }
 
 // ClientCloser is used to close Client from further use
@@ -135,28 +136,40 @@ func NewClient(addr string, namespace string, handler interface{}) (ClientCloser
 				}
 			}
 
+			var ctx context.Context
+			if hasCtx == 1 {
+				ctx = args[0].Interface().(context.Context)
+			}
+
 			var retVal reflect.Value
-			var chCtor func() func([]byte, bool)
+
+			// if the function returns a channel, we need to provide a sink for the
+			// messages
+			var chCtor func() (context.Context, func([]byte, bool))
+
 			if retCh {
 				retVal = reflect.Zero(ftyp.Out(valOut))
-				chCtor = func() func([]byte, bool) {
+
+				chCtor = func() (context.Context, func([]byte, bool)) {
 					// unpack chan type to make sure it's reflect.BothDir
 					ctyp := reflect.ChanOf(reflect.BothDir, ftyp.Out(valOut).Elem())
 					ch := reflect.MakeChan(ctyp, 0) // todo: buffer?
 					retVal = ch.Convert(ftyp.Out(valOut))
 
-					return func(result []byte, ok bool) {
+					return ctx, func(result []byte, ok bool) {
 						if !ok {
+							// remote channel closed, close ours too
 							ch.Close()
 							return
 						}
+
 						val := reflect.New(ftyp.Out(valOut).Elem())
 						if err := json.Unmarshal(result, val.Interface()); err != nil {
 							log.Errorf("error unmarshaling chan response: %s", err)
 							return
 						}
 
-						ch.Send(val.Elem()) // todo: select on ctx
+						ch.Send(val.Elem()) // todo: select on ctx is probably a good idea
 					}
 				}
 			}
@@ -178,8 +191,8 @@ func NewClient(addr string, namespace string, handler interface{}) (ClientCloser
 			var ctxDone <-chan struct{}
 			var resp clientResponse
 
-			if hasCtx == 1 {
-				ctxDone = args[0].Interface().(context.Context).Done()
+			if ctx != nil {
+				ctxDone = ctx.Done()
 			}
 
 			// wait for response, handle context cancellation

lib/jsonrpc/rpc_test.go

@@ -3,6 +3,7 @@ package jsonrpc
 import (
 	"context"
 	"errors"
+	"fmt"
 	"net/http/httptest"
 	"strconv"
 	"strings"
@@ -241,11 +242,12 @@ func TestCtx(t *testing.T) {
 }
 
 type UnUnmarshalable int
+
 func (*UnUnmarshalable) UnmarshalJSON([]byte) error {
 	return errors.New("nope")
 }
 
-type UnUnmarshalableHandler struct {}
+type UnUnmarshalableHandler struct{}
 
 func (*UnUnmarshalableHandler) GetUnUnmarshalableStuff() (UnUnmarshalable, error) {
 	return UnUnmarshalable(5), nil
@@ -274,7 +276,7 @@ type ChanHandler struct {
 	wait chan struct{}
 }
 
-func (h *ChanHandler) Sub(ctx context.Context, i int) (<-chan int, error) {
+func (h *ChanHandler) Sub(ctx context.Context, i int, eq int) (<-chan int, error) {
 	out := make(chan int)
 
 	go func() {
@@ -282,11 +284,23 @@ func (h *ChanHandler) Sub(ctx context.Context, i int) (<-chan int, error) {
 		var n int
 
 		for {
-			<-h.wait
+			select {
+			case <-ctx.Done():
+				fmt.Println("ctxdone1")
+				return
+			case <-h.wait:
+			}
+
 			n += i
 
+			if n == eq {
+				fmt.Println("eq")
+				return
+			}
+
 			select {
 			case <-ctx.Done():
+				fmt.Println("ctxdone2")
 				return
 			case out <- n:
 			}
@@ -298,7 +312,7 @@ func (h *ChanHandler) Sub(ctx context.Context, i int) (<-chan int, error) {
 
 func TestChan(t *testing.T) {
 	var client struct {
-		Sub func(context.Context, int) (<-chan int, error)
+		Sub func(context.Context, int, int) (<-chan int, error)
 	}
 
 	serverHandler := &ChanHandler{
@@ -319,13 +333,47 @@ func TestChan(t *testing.T) {
 	serverHandler.wait <- struct{}{}
 
 	ctx, cancel := context.WithCancel(context.Background())
-	defer cancel()
 
+	// sub
 
-
-	sub, err := client.Sub(ctx, 2)
+	sub, err := client.Sub(ctx, 2, -1)
 	require.NoError(t, err)
 
+	// recv one
 
 	require.Equal(t, 2, <-sub)
+
+	// recv many (order)
+
+	serverHandler.wait <- struct{}{}
+	serverHandler.wait <- struct{}{}
+	serverHandler.wait <- struct{}{}
+
+	require.Equal(t, 4, <-sub)
+	require.Equal(t, 6, <-sub)
+	require.Equal(t, 8, <-sub)
+
+	// close (through ctx)
+	cancel()
+
+	_, ok := <-sub
+	require.Equal(t, false, ok)
+
+	// sub (again)
+
+	serverHandler.wait <- struct{}{}
+
+	ctx, cancel = context.WithCancel(context.Background())
+	defer cancel()
+
+	sub, err = client.Sub(ctx, 3, 6)
+	require.NoError(t, err)
+
+	require.Equal(t, 3, <-sub)
+
+	// close (remote)
+	serverHandler.wait <- struct{}{}
+	_, ok = <-sub
+	require.Equal(t, false, ok)
+
 }

lib/jsonrpc/websocket.go

@@ -145,8 +145,16 @@ func handleWsConn(ctx context.Context, conn *websocket.Conn, handler handlers, r
 					caseToId[chosen-1] = caseToId[n-1]
 				}
 
+				id := caseToId[chosen-1]
 				cases = cases[:n]
 				caseToId = caseToId[:n-1]
+
+				sendReq(request{
+					Jsonrpc: "2.0",
+					ID:      nil, // notification
+					Method:  chClose,
+					Params:  []param{{v: reflect.ValueOf(id)}},
+				})
 				continue
 			}
 
@@ -199,6 +207,16 @@ func handleWsConn(ctx context.Context, conn *websocket.Conn, handler handlers, r
 		}
 	}()
 
+	handleCtxAsync := func(actx context.Context, id int64) {
+		<-actx.Done()
+
+		sendReq(request{
+			Jsonrpc: "2.0",
+			Method:  wsCancel,
+			Params:  []param{{v: reflect.ValueOf(id)}},
+		})
+	}
+
 	// cancelCtx is a built-in rpc which handles context cancellation over rpc
 	cancelCtx := func(req frame) {
 		if req.ID != nil {
@@ -259,7 +277,9 @@ func handleWsConn(ctx context.Context, conn *websocket.Conn, handler handlers, r
 						continue
 					}
 
-					chanHandlers[chid] = req.retCh()
+					var chanCtx context.Context
+					chanCtx, chanHandlers[chid] = req.retCh()
+					go handleCtxAsync(chanCtx, *frame.ID)
 				}
 
 				req.ready <- clientResponse{
@@ -285,6 +305,22 @@ func handleWsConn(ctx context.Context, conn *websocket.Conn, handler handlers, r
 				}
 
 				hnd(frame.Params[1].data, true)
+			case chClose:
+				var chid uint64
+				if err := json.Unmarshal(frame.Params[0].data, &chid); err != nil {
+					log.Error("failed to unmarshal channel id in xrpc.ch.val: %s", err)
+					continue
+				}
+
+				hnd, ok := chanHandlers[chid]
+				if !ok {
+					log.Errorf("xrpc.ch.val: handler %d not found", chid)
+					continue
+				}
+
+				delete(chanHandlers, chid)
+
+				hnd(nil, false)
 			default: // Remote call
 				req := request{
 					Jsonrpc: frame.Jsonrpc,
@@ -316,9 +352,8 @@ func handleWsConn(ctx context.Context, conn *websocket.Conn, handler handlers, r
 
 						if !keepctx {
 							cf()
+							delete(handling, *frame.ID)
 						}
-
-						delete(handling, *frame.ID)
 					}
 				}