Repository: creachadair/jrpc2
Pull request #66 (2021-12-07 02:40:31)

server: update context plumbing
Update the Accepter interface to take a context argument.
This is a breaking change to the interface, but NetAccepter works
without change, so most existing use is not affected.

Rework Loop as LoopContext, which takes a context argument to
govern the lifecycle of the loop itself. Retain Loop as a wrapper
for LoopContext with a background context.

Update and clean up tests.


Patch:
diff --git a/server/loop.go b/server/loop.go
index e51a75e..7cdf61c 100644
--- a/server/loop.go
+++ b/server/loop.go
@@ -4,6 +4,7 @@
 package server
 
 import (
+	"context"
 	"net"
 	"sync"
 
@@ -39,8 +40,9 @@ func (static) Finish(jrpc2.Assigner, jrpc2.ServerStatus) {}
 // An Accepter obtains client connections from an external source and
 // constructs channels from them.
 type Accepter interface {
-	// Accept accepts a connection and returns a new channel for it.
-	Accept() (channel.Channel, error)
+	// Accept blocks until a connection is available, or until ctx ends.
+	// If a connection is found, Accept  returns a new channel for it.
+	Accept(ctx context.Context) (channel.Channel, error)
 }
 
 // NetAccepter adapts a net.Listener to the Accepter interface, using f as the
@@ -54,7 +56,20 @@ type netAccepter struct {
 	newChannel channel.Framing
 }
 
-func (n netAccepter) Accept() (channel.Channel, error) {
+func (n netAccepter) Accept(ctx context.Context) (channel.Channel, error) {
+	// A net.Listener does not obey a context, so simulate it by closing the
+	// listener if ctx ends.
+	ok := make(chan struct{})
+	defer close(ok)
+	go func() {
+		select {
+		case <-ctx.Done():
+			n.Listener.Close()
+		case <-ok:
+			return
+		}
+	}()
+
 	conn, err := n.Listener.Accept()
 	if err != nil {
 		return nil, err
@@ -62,14 +77,19 @@ func (n netAccepter) Accept() (channel.Channel, error) {
 	return n.newChannel(conn, conn), nil
 }
 
-// Loop obtains connections from lst and starts a server for each using a
-// service instance returned by newService and the given options. Each server
+// Loop calls LoopContext with the given arguments and a background context.
+// See LoopContext for details.
+func Loop(lst Accepter, newService func() Service, opts *LoopOptions) error {
+	return LoopContext(context.Background(), lst, newService, opts)
+}
+
+// LoopContext obtains connections from lst and starts a server for each using
+// a service instance returned by newService and the given options. Each server
 // runs in a new goroutine.
 //
-// If the listener reports an error, the loop will terminate and that error
-// will be reported to the caller of Loop once any active servers have
-// returned.
-func Loop(lst Accepter, newService func() Service, opts *LoopOptions) error {
+// If lst reports an error, the loop will terminate and that error will be
+// reported to the caller of Loop once any active servers have returned.
+func LoopContext(ctx context.Context, lst Accepter, newService func() Service, opts *LoopOptions) error {
 	serverOpts := opts.serverOpts()
 	log := func(string, ...interface{}) {}
 	if serverOpts != nil && serverOpts.Logger != nil {
@@ -78,7 +98,7 @@ func Loop(lst Accepter, newService func() Service, opts *LoopOptions) error {
 
 	var wg sync.WaitGroup
 	for {
-		ch, err := lst.Accept()
+		ch, err := lst.Accept(ctx)
 		if err != nil {
 			if channel.IsErrClosing(err) {
 				err = nil


Test patch:
diff --git a/server/loop_test.go b/server/loop_test.go
index 8ac68b6..30341e7 100644
--- a/server/loop_test.go
+++ b/server/loop_test.go
@@ -1,6 +1,6 @@
 // Copyright (C) 2017 Michael J. Fromberger. All Rights Reserved.
 
-package server
+package server_test
 
 import (
 	"context"
@@ -13,13 +13,14 @@ import (
 	"github.com/creachadair/jrpc2"
 	"github.com/creachadair/jrpc2/channel"
 	"github.com/creachadair/jrpc2/handler"
+	"github.com/creachadair/jrpc2/server"
 	"github.com/fortytw2/leaktest"
 )
 
 var newChan = channel.Line
 
 // A static test service that returns the same thing each time.
-var testService = Static(handler.Map{
+var testStatic = server.Static(handler.Map{
 	"Test": handler.New(func(context.Context) (string, error) {
 		return "OK", nil
 	}),
@@ -32,8 +33,8 @@ type testSession struct {
 	nCall int
 }
 
-func newTestSession(t *testing.T) func() Service {
-	return func() Service { t.Helper(); return &testSession{t: t} }
+func newTestSession(t *testing.T) func() server.Service {
+	return func() server.Service { t.Helper(); return &testSession{t: t} }
 }
 
 func (t *testSession) Assigner() (jrpc2.Assigner, error) {
@@ -84,20 +85,18 @@ func mustDial(t *testing.T, addr string) *jrpc2.Client {
 	return jrpc2.NewClient(newChan(conn, conn), nil)
 }
 
-func mustServe(t *testing.T, lst net.Listener, newService func() Service) <-chan struct{} {
+func mustServe(t *testing.T, ctx context.Context, lst net.Listener, newService func() server.Service) <-chan error {
 	t.Helper()
 
-	sc := make(chan struct{})
+	acc := server.NetAccepter(lst, newChan)
+	errc := make(chan error, 1)
 	go func() {
-		defer close(sc)
+		defer close(errc)
 		// Start a server loop to accept connections from the clients. This should
 		// exit cleanly once all the clients have finished and the listener closes.
-		lst := NetAccepter(lst, newChan)
-		if err := Loop(lst, newService, nil); err != nil {
-			t.Errorf("Loop: unexpected failure: %v", err)
-		}
+		errc <- server.LoopContext(ctx, acc, newService, nil)
 	}()
-	return sc
+	return errc
 }
 
 // Test that sequential clients against the same server work sanely.
@@ -106,7 +105,7 @@ func TestSeq(t *testing.T) {
 
 	lst := mustListen(t)
 	addr := lst.Addr().String()
-	sc := mustServe(t, lst, testService)
+	errc := mustServe(t, context.Background(), lst, testStatic)
 
 	for i := 0; i < 5; i++ {
 		cli := mustDial(t, addr)
@@ -119,7 +118,33 @@ func TestSeq(t *testing.T) {
 		cli.Close()
 	}
 	lst.Close()
-	<-sc
+	if err := <-errc; err != nil {
+		t.Errorf("Server exit failed: %v", err)
+	}
+}
+
+// Test that context plumbing works properly.
+func TestLoopContext_cancellation(t *testing.T) {
+	defer leaktest.Check(t)()
+
+	ctx, cancel := context.WithCancel(context.Background())
+	lst := server.NetAccepter(mustListen(t), channel.Line)
+
+	errc := make(chan error, 1)
+	go func() {
+		defer close(errc)
+		errc <- server.LoopContext(ctx, lst, testStatic, nil)
+	}()
+
+	time.AfterFunc(50*time.Millisecond, cancel)
+	select {
+	case err := <-errc:
+		if err != nil {
+			t.Errorf("Loop exit reported error: %v", err)
+		}
+	case <-time.After(1 * time.Second):
+		t.Fatalf("Loop did not exit in a timely manner after cancellation")
+	}
 }
 
 // Test that concurrent clients against the same server work sanely.
@@ -128,9 +153,9 @@ func TestLoop(t *testing.T) {
 
 	tests := []struct {
 		desc string
-		cons func() Service
+		cons func() server.Service
 	}{
-		{"StaticService", testService},
+		{"StaticService", testStatic},
 		{"SessionStateService", newTestSession(t)},
 	}
 	const numClients = 5
@@ -140,7 +165,7 @@ func TestLoop(t *testing.T) {
 		t.Run(test.desc, func(t *testing.T) {
 			lst := mustListen(t)
 			addr := lst.Addr().String()
-			sc := mustServe(t, lst, test.cons)
+			errc := mustServe(t, context.Background(), lst, test.cons)
 
 			// Start a bunch of clients, each of which will dial the server and make
 			// some calls at random intervals to tickle the race detector.
@@ -169,7 +194,9 @@ func TestLoop(t *testing.T) {
 			// the service loop will stop.
 			wg.Wait()
 			lst.Close()
-			<-sc
+			if err := <-errc; err != nil {
+				t.Errorf("Server exit failed: %v", err)
+			}
 		})
 	}
 }