Test RPC server external context cancellation

rpc/rpc_test.go

@@ -0,0 +1,119 @@
+package rpc
+
+import (
+	"context"
+	"net"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"capnproto.org/go/capnp/v3"
+	"capnproto.org/go/capnp/v3/rpc/internal/testcapnp"
+)
+
+func TestConnection_BaseContext(t *testing.T) {
+	t.Parallel()
+
+	t.Run("background context", func(t *testing.T) {
+		client, server := net.Pipe()
+		doneCh := make(chan struct{}, 1)
+
+		go func() {
+			bootstrapClient := testcapnp.StreamTest_ServerToClient(slowStreamTestServer{})
+			conn := NewConn(NewStreamTransport(server), &Options{
+				BootstrapClient: capnp.Client(bootstrapClient),
+			})
+			defer conn.Close()
+
+			<-conn.Done()
+			close(doneCh)
+		}()
+
+		func() {
+			conn := NewConn(NewStreamTransport(client), nil)
+			defer conn.Close()
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			client := testcapnp.StreamTest(conn.Bootstrap(ctx))
+			defer client.Release()
+
+			err := client.Push(ctx, func(st testcapnp.StreamTest_push_Params) error {
+				return st.SetData(make([]byte, 1))
+			})
+
+			require.NoError(t, err)
+		}()
+
+		<-doneCh
+	})
+
+	t.Run("external context", func(t *testing.T) {
+		client, server := net.Pipe()
+
+		ctx, cancel := context.WithCancel(context.Background())
+		defer cancel()
+
+		clientConnectedCh := make(chan struct{}, 1)
+		contextCancelledCh := make(chan struct{}, 1)
+		wg := &sync.WaitGroup{}
+		wg.Add(2)
+
+		go func() {
+			defer wg.Done()
+
+			ctx, cancel := context.WithCancel(ctx)
+			defer cancel()
+
+			bootstrapClient := testcapnp.StreamTest_ServerToClient(slowStreamTestServer{})
+			conn := NewConn(NewStreamTransport(server), &Options{
+				BootstrapClient: capnp.Client(bootstrapClient),
+				BaseContext: func() context.Context {
+					return ctx
+				},
+			})
+			defer conn.Close()
+
+			select {
+			case <-clientConnectedCh:
+				cancel()
+				contextCancelledCh <- struct{}{}
+			case <-conn.Done():
+				t.Failed()
+			}
+
+			// Connection should close when external context is cancelled
+			<-conn.Done()
+		}()
+
+		go func() {
+			defer wg.Done()
+
+			conn := NewConn(NewStreamTransport(client), nil)
+			defer conn.Close()
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			client := testcapnp.StreamTest(conn.Bootstrap(ctx))
+			defer client.Release()
+
+			if err := client.Resolve(ctx); err != nil {
+				require.NoError(t, err)
+			}
+
+			clientConnectedCh <- struct{}{}
+			<-contextCancelledCh
+
+			err := client.Push(ctx, func(st testcapnp.StreamTest_push_Params) error {
+				return st.SetData(make([]byte, 1))
+			})
+
+			require.NoError(t, err)
+		}()
+
+		wg.Wait()
+	})
+}