feat: support context for server acceptor (#97)

jjeffcaii · web-flow · commit 318cdc638b4e · 2021-01-26T22:21:52.000+08:00
diff --git a/README.md b/README.md
@@ -45,7 +45,7 @@ import (
 
 func main() {
 	err := rsocket.Receive().
-		Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+		Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 			// bind responder
 			return rsocket.NewAbstractSocket(
 				rsocket.RequestResponse(func(msg payload.Payload) mono.Mono {
diff --git a/balancer/group_example_test.go b/balancer/group_example_test.go
@@ -21,7 +21,7 @@ func ExampleNewGroup() {
 	// Create a broker with resume.
 	err := rsocket.Receive().
 		Resume(rsocket.WithServerResumeSessionDuration(10 * time.Second)).
-		Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+		Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 			// Register service using Setup Metadata as service ID.
 			if serviceID, ok := setup.MetadataUTF8(); ok {
 				_ = group.Get(serviceID).Put(sendingSocket)
diff --git a/balancer/round_robin_test.go b/balancer/round_robin_test.go
@@ -20,7 +20,7 @@ import (
 
 func startServer(ctx context.Context, port int, counter *sync.Map) {
 	_ = rsocket.Receive().
-		Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+		Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 			return rsocket.NewAbstractSocket(
 				rsocket.RequestResponse(func(msg payload.Payload) mono.Mono {
 					cur, _ := counter.LoadOrStore(port, atomic.NewInt32(0))
diff --git a/client.go b/client.go
@@ -30,7 +30,7 @@ type Client interface {
 }
 
 // ClientSocketAcceptor is alias for RSocket handler function.
-type ClientSocketAcceptor = func(socket RSocket) RSocket
+type ClientSocketAcceptor = func(ctx context.Context, socket RSocket) RSocket
 
 // ClientStarter can be used to start a client.
 type ClientStarter interface {
@@ -192,13 +192,17 @@ func (cb *clientBuilder) Transport(t transport.ClientTransporter) ClientStarter
 }
 
 func (cb *clientBuilder) Start(ctx context.Context) (client Client, err error) {
+	if ctx == nil {
+		panic("context cannot be nil!")
+	}
 	// create a blank socket.
 	err = fragmentation.IsValidFragment(cb.fragment)
 	if err != nil {
 		return
 	}
 
 	conn := socket.NewClientDuplexConnection(
+		ctx,
 		cb.reqSche,
 		cb.resSche,
 		cb.fragment,
@@ -213,7 +217,7 @@ func (cb *clientBuilder) Start(ctx context.Context) (client Client, err error) {
 		cs = socket.NewClient(cb.tpGen, conn)
 	}
 	if cb.acceptor != nil {
-		conn.SetResponder(cb.acceptor(cs))
+		conn.SetResponder(cb.acceptor(ctx, cs))
 	} else {
 		conn.SetResponder(_noopSocket)
 	}
diff --git a/cmd/rsocket-cli/runner.go b/cmd/rsocket-cli/runner.go
@@ -182,7 +182,7 @@ func (r *Runner) runServerMode(ctx context.Context) error {
 	go func() {
 		sendingPayloads := r.createPayload()
 		ch <- sb.
-			Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+			Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 				var options []rsocket.OptAbstractSocket
 				options = append(options, rsocket.RequestStream(func(message payload.Payload) flux.Flux {
 					r.showPayload(message)
diff --git a/examples/echo/echo.go b/examples/echo/echo.go
@@ -31,7 +31,7 @@ func main() {
 		OnStart(func() {
 			log.Println("server start success!")
 		}).
-		Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+		Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 			//log.Println("SETUP BEGIN:----------------")
 			//log.Println("maxLifeTime:", setup.MaxLifetime())
 			//log.Println("keepaliveInterval:", setup.TimeBetweenKeepalive())
diff --git a/examples/echo_bench/echo_bench.go b/examples/echo_bench/echo_bench.go
@@ -91,7 +91,7 @@ func createClient(mtu int) (rsocket.Client, error) {
 			log.Println("*** disconnected ***", common.CountBorrowed())
 		}).
 		SetupPayload(payload.NewString("你好", "世界")).
-		Acceptor(func(socket rsocket.RSocket) rsocket.RSocket {
+		Acceptor(func(ctx context.Context, socket rsocket.RSocket) rsocket.RSocket {
 			return rsocket.NewAbstractSocket(
 				rsocket.RequestResponse(func(p payload.Payload) mono.Mono {
 					log.Println("receive request from server:", p)
diff --git a/examples/fibonacci/main.go b/examples/fibonacci/main.go
@@ -71,7 +71,7 @@ func server(readyCh chan struct{}) {
 			// close the channel to singal that the server is ready
 			close(readyCh)
 		}).
-		Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+		Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 			// register a new request stream handler
 			return rsocket.NewAbstractSocket(requestStreamHandler), nil
 		}).
diff --git a/examples/word_counter/main.go b/examples/word_counter/main.go
@@ -49,7 +49,7 @@ func server(readyCh chan struct{}) {
 			// close the channel to signal that the server is ready
 			close(readyCh)
 		}).
-		Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+		Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 			// register a new request channel handler
 			return rsocket.NewAbstractSocket(requestChannelHandler), nil
 		}).
diff --git a/go.mod b/go.mod
@@ -6,7 +6,7 @@ require (
 	github.com/golang/mock v1.4.4
 	github.com/google/uuid v1.1.2
 	github.com/gorilla/websocket v1.4.2
-	github.com/jjeffcaii/reactor-go v0.5.0
+	github.com/jjeffcaii/reactor-go v0.5.1
 	github.com/pkg/errors v0.9.1
 	github.com/pkg/profile v1.5.0
 	github.com/stretchr/testify v1.6.1
diff --git a/go.sum b/go.sum
@@ -12,6 +12,8 @@ github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0U
 github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 github.com/jjeffcaii/reactor-go v0.5.0 h1:Y8VVp31JGij/ilMh8bXyNEo8k2SVSmc4Cc+mzPGL07w=
 github.com/jjeffcaii/reactor-go v0.5.0/go.mod h1:qYN34C2UANFOtDeUGvhxlExLpFMDbvSrphX3Gb3H6S8=
+github.com/jjeffcaii/reactor-go v0.5.1 h1:YG2SUE/Pag2ChZOcJ+Twfb+D10lNOavoYly+xUUbA7w=
+github.com/jjeffcaii/reactor-go v0.5.1/go.mod h1:qYN34C2UANFOtDeUGvhxlExLpFMDbvSrphX3Gb3H6S8=
 github.com/panjf2000/ants/v2 v2.4.3 h1:wHghL17YKFanB62QjPQ9o+DuM4q7WrQ7zAhoX8+eBXU=
 github.com/panjf2000/ants/v2 v2.4.3/go.mod h1:f6F0NZVFsGCp5A7QW/Zj/m92atWwOkY0OIhFxRNFr4A=
 github.com/pkg/errors v0.9.1 h1:FEBLx1zS214owpjy7qsBeixbURkuhQAwrK5UwLGTwt4=
diff --git a/internal/socket/base_socket_test.go b/internal/socket/base_socket_test.go
@@ -24,7 +24,7 @@ func TestBaseSocket(t *testing.T) {
 	conn.EXPECT().Read().Return(nil, io.EOF).AnyTimes()
 	conn.EXPECT().SetDeadline(gomock.Any()).AnyTimes()
 
-	duplex := socket.NewClientDuplexConnection(nil, nil, fragmentation.MaxFragment, 90*time.Second)
+	duplex := socket.NewClientDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, 90*time.Second)
 	duplex.SetTransport(tp)
 
 	go func() {
diff --git a/internal/socket/duplex.go b/internal/socket/duplex.go
@@ -47,6 +47,7 @@ func mustExecute(sc scheduler.Scheduler, handler func()) {
 
 // DuplexConnection represents a socket of RSocket which can be a requester or a responder.
 type DuplexConnection struct {
+	ctx            context.Context
 	reqSche        scheduler.Scheduler
 	resSche        scheduler.Scheduler
 	destroyReqSche bool
@@ -504,7 +505,7 @@ func (dc *DuplexConnection) RequestChannel(sending flux.Flux) (ret flux.Flux) {
 				rcv:          receiving,
 				result:       sendResult,
 			}
-			sending.SubscribeOn(dc.reqSche).SubscribeWith(context.Background(), sub)
+			sending.SubscribeOn(dc.reqSche).SubscribeWith(dc.ctx, sub)
 		})
 	return ret
 }
@@ -551,11 +552,11 @@ func (dc *DuplexConnection) respondRequestResponse(receiving fragmentation.Heade
 	// async subscribe publisher
 	sub := borrowRequestResponseSubscriber(dc, sid, receiving)
 	if mono.IsSubscribeAsync(sending) {
-		sending.SubscribeWith(context.Background(), sub)
+		sending.SubscribeWith(dc.ctx, sub)
 		return nil
 	}
 	mustExecute(dc.resSche, func() {
-		sending.SubscribeWith(context.Background(), sub)
+		sending.SubscribeWith(dc.ctx, sub)
 	})
 	return nil
 }
@@ -659,7 +660,7 @@ func (dc *DuplexConnection) respondRequestChannel(req fragmentation.HeaderAndPay
 	}
 
 	mustExecute(dc.reqSche, func() {
-		sending.SubscribeWith(context.Background(), sub)
+		sending.SubscribeWith(dc.ctx, sub)
 	})
 
 	<-subscribed
@@ -761,7 +762,7 @@ func (dc *DuplexConnection) respondRequestStream(receiving fragmentation.HeaderA
 
 	// async subscribe publisher
 	sub := borrowRequestStreamSubscriber(receiving, dc, sid, n)
-	sending.SubscribeOn(dc.resSche).SubscribeWith(context.Background(), sub)
+	sending.SubscribeOn(dc.resSche).SubscribeWith(dc.ctx, sub)
 	return nil
 }
 
@@ -1344,16 +1345,16 @@ func IsSocketClosedError(err error) bool {
 }
 
 // NewServerDuplexConnection creates a new server-side DuplexConnection.
-func NewServerDuplexConnection(reqSche, resSche scheduler.Scheduler, mtu int, leases lease.Factory) *DuplexConnection {
-	return newDuplexConnection(reqSche, resSche, mtu, nil, &serverStreamIDs{}, leases)
+func NewServerDuplexConnection(ctx context.Context, reqSche, resSche scheduler.Scheduler, mtu int, leases lease.Factory) *DuplexConnection {
+	return newDuplexConnection(ctx, reqSche, resSche, mtu, nil, &serverStreamIDs{}, leases)
 }
 
 // NewClientDuplexConnection creates a new client-side DuplexConnection.
-func NewClientDuplexConnection(reqSche, resSche scheduler.Scheduler, mtu int, keepaliveInterval time.Duration) *DuplexConnection {
-	return newDuplexConnection(reqSche, resSche, mtu, NewKeepaliver(keepaliveInterval), &clientStreamIDs{}, nil)
+func NewClientDuplexConnection(ctx context.Context, reqSche, resSche scheduler.Scheduler, mtu int, keepaliveInterval time.Duration) *DuplexConnection {
+	return newDuplexConnection(ctx, reqSche, resSche, mtu, NewKeepaliver(keepaliveInterval), &clientStreamIDs{}, nil)
 }
 
-func newDuplexConnection(reqSche, resSche scheduler.Scheduler, mtu int, ka *Keepaliver, sids StreamID, leases lease.Factory) *DuplexConnection {
+func newDuplexConnection(ctx context.Context, reqSche, resSche scheduler.Scheduler, mtu int, ka *Keepaliver, sids StreamID, leases lease.Factory) *DuplexConnection {
 	destroyReqSche := reqSche == nil
 	if destroyReqSche {
 		reqSche = scheduler.NewElastic(misc.MaxInt(runtime.NumCPU()<<8, _minRequestSchedulerSize))
@@ -1363,6 +1364,7 @@ func newDuplexConnection(reqSche, resSche scheduler.Scheduler, mtu int, ka *Keep
 	}
 
 	c := &DuplexConnection{
+		ctx:            ctx,
 		reqSche:        reqSche,
 		resSche:        resSche,
 		destroyReqSche: destroyReqSche,
diff --git a/internal/socket/resumable_client_socket_test.go b/internal/socket/resumable_client_socket_test.go
@@ -49,7 +49,7 @@ func TestNewResumableClientSocket(t *testing.T) {
 	}).AnyTimes()
 	conn.EXPECT().SetDeadline(gomock.Any()).AnyTimes()
 
-	ds := socket.NewClientDuplexConnection(nil, nil, fragmentation.MaxFragment, 90*time.Second)
+	ds := socket.NewClientDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, 90*time.Second)
 
 	rcs := socket.NewResumableClientSocket(func(ctx context.Context) (*transport.Transport, error) {
 		return tp, nil
@@ -109,7 +109,7 @@ func TestNewResumableClientSocket(t *testing.T) {
 }
 
 func TestResumeClientSocket_Setup_Broken(t *testing.T) {
-	c := socket.NewClientDuplexConnection(nil, nil, fragmentation.MaxFragment, 90*time.Second)
+	c := socket.NewClientDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, 90*time.Second)
 	s := socket.NewResumableClientSocket(func(ctx context.Context) (*transport.Transport, error) {
 		return nil, fakeErr
 	}, c)
@@ -122,7 +122,7 @@ func TestResumeClientSocket_Setup(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	defer ctrl.Finish()
 
-	ds := socket.NewClientDuplexConnection(nil, nil, fragmentation.MaxFragment, 90*time.Second)
+	ds := socket.NewClientDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, 90*time.Second)
 
 	readChanChan := make(chan chan core.BufferedFrame, 64)
 
diff --git a/internal/socket/resumable_server_socket_test.go b/internal/socket/resumable_server_socket_test.go
@@ -54,7 +54,7 @@ func TestResumableServerSocket_Start(t *testing.T) {
 
 	close(readChan)
 
-	c := socket.NewServerDuplexConnection(nil, nil, fragmentation.MaxFragment, nil)
+	c := socket.NewServerDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, nil)
 	ss := socket.NewResumableServerSocket(c, fakeToken)
 
 	ss.SetResponder(fakeResponder)
diff --git a/internal/socket/simple_client_socket_test.go b/internal/socket/simple_client_socket_test.go
@@ -19,7 +19,7 @@ import (
 )
 
 func TestNewClientWithBrokenTransporter(t *testing.T) {
-	ds := socket.NewClientDuplexConnection(nil, nil, fragmentation.MaxFragment, 90*time.Second)
+	ds := socket.NewClientDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, 90*time.Second)
 	// Must failed transporter
 	transporter := func(ctx context.Context) (*transport.Transport, error) {
 		return nil, fakeErr
@@ -49,7 +49,7 @@ func TestNewClient(t *testing.T) {
 	}).AnyTimes()
 	conn.EXPECT().SetDeadline(gomock.Any()).AnyTimes()
 
-	ds := socket.NewClientDuplexConnection(nil, nil, fragmentation.MaxFragment, 90*time.Second)
+	ds := socket.NewClientDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, 90*time.Second)
 	cli := socket.NewClient(func(ctx context.Context) (*transport.Transport, error) {
 		return tp, nil
 	}, ds)
@@ -120,7 +120,7 @@ func TestLease(t *testing.T) {
 	}).AnyTimes()
 	conn.EXPECT().SetDeadline(gomock.Any()).AnyTimes()
 
-	ds := socket.NewClientDuplexConnection(nil, nil, fragmentation.MaxFragment, 90*time.Second)
+	ds := socket.NewClientDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, 90*time.Second)
 	cli := socket.NewClient(func(ctx context.Context) (*transport.Transport, error) {
 		return tp, nil
 	}, ds)
diff --git a/internal/socket/simple_server_socket_test.go b/internal/socket/simple_server_socket_test.go
@@ -55,7 +55,7 @@ func TestSimpleServerSocket_Start(t *testing.T) {
 
 	close(readChan)
 
-	c := socket.NewServerDuplexConnection(nil, nil, fragmentation.MaxFragment, nil)
+	c := socket.NewServerDuplexConnection(context.Background(), nil, nil, fragmentation.MaxFragment, nil)
 	ss := socket.NewSimpleServerSocket(c)
 	ss.SetResponder(fakeResponder)
 	ss.SetTransport(tp)
diff --git a/internal/stopwatch/Stopwatch.go b/internal/stopwatch/Stopwatch.go
diff --git a/rsocket.go b/rsocket.go
@@ -1,6 +1,8 @@
 package rsocket
 
 import (
+	"context"
+
 	"github.com/rsocket/rsocket-go/core"
 	"github.com/rsocket/rsocket-go/internal/socket"
 	"github.com/rsocket/rsocket-go/payload"
@@ -41,7 +43,7 @@ type (
 
 type (
 	// ServerAcceptor is alias for server acceptor.
-	ServerAcceptor = func(setup payload.SetupPayload, sendingSocket CloseableRSocket) (RSocket, error)
+	ServerAcceptor = func(ctx context.Context, setup payload.SetupPayload, socket CloseableRSocket) (RSocket, error)
 
 	// RSocket is a contract providing different interaction models for RSocket protocol.
 	RSocket interface {
diff --git a/rsocket_example_test.go b/rsocket_example_test.go
@@ -17,7 +17,7 @@ import (
 func Example() {
 	// Serve a server
 	err := rsocket.Receive().
-		Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+		Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 			return rsocket.NewAbstractSocket(
 				rsocket.RequestResponse(func(msg payload.Payload) mono.Mono {
 					log.Println("incoming request:", msg)
@@ -52,7 +52,7 @@ func ExampleReceive() {
 	err := rsocket.Receive().
 		Resume(rsocket.WithServerResumeSessionDuration(30 * time.Second)).
 		Fragment(65535).
-		Acceptor(func(setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
+		Acceptor(func(ctx context.Context, setup payload.SetupPayload, sendingSocket rsocket.CloseableRSocket) (rsocket.RSocket, error) {
 			// Handle close.
 			sendingSocket.OnClose(func(err error) {
 				log.Println("sending socket is closed")
@@ -101,7 +101,7 @@ func ExampleConnect() {
 		Lease().  // Enable LEASE.
 		Fragment(4096).
 		SetupPayload(payload.NewString("Hello", "World")).
-		Acceptor(func(socket rsocket.RSocket) rsocket.RSocket {
+		Acceptor(func(ctx context.Context, socket rsocket.RSocket) rsocket.RSocket {
 			return rsocket.NewAbstractSocket(
 				rsocket.RequestResponse(func(msg payload.Payload) mono.Mono {
 					return mono.Just(payload.NewString("Pong", time.Now().String()))
diff --git a/rsocket_test.go b/rsocket_test.go
diff --git a/server.go b/server.go
