From c67d45b065f73b306dac72148ae67a9f0ff7c272 Mon Sep 17 00:00:00 2001 From: Dzmitry Mikhalapau Date: Wed, 27 Nov 2024 23:46:01 +0500 Subject: [PATCH 1/2] Provide an option to specify base context for RPC connection Signed-off-by: Dzmitry Mikhalapau --- rpc/rpc.go | 21 ++++++++++++++++----- 1 file changed, 16 insertions(+), 5 deletions(-) diff --git a/rpc/rpc.go b/rpc/rpc.go index a557ba8b..fedd76d6 100644 --- a/rpc/rpc.go +++ b/rpc/rpc.go @@ -88,10 +88,11 @@ type Conn struct { bootstrap capnp.Client er errReporter abortTimeout time.Duration + baseContext func() context.Context // bgctx is a Context that is canceled when shutdown starts. Note - // that it's parent is context.Background(), so we can rely on this - // being the *only* time it will be canceled. + // that if baseContext is not provided. it's parent is context.Background(), + // so we can rely on this being the *only* time it will be canceled. bgctx context.Context // tasks block shutdown. @@ -202,6 +203,11 @@ type Options struct { // by Dial or Accept on the Network itself; application code should not // set this. Network Network + + // BaseContext is an optional funcation that returns a base context + // for any incoming connection. If ommitted, the context.Background() + // will be used instead. + BaseContext func() context.Context } // Logger is used for logging by the RPC system. Each method logs @@ -231,8 +237,9 @@ type Logger interface { // requests from the transport. func NewConn(t Transport, opts *Options) *Conn { c := &Conn{ - transport: t, - closed: make(chan struct{}), + transport: t, + baseContext: context.Background, + closed: make(chan struct{}), } sender := spsc.New[asyncSend]() @@ -248,6 +255,10 @@ func NewConn(t Transport, opts *Options) *Conn { c.abortTimeout = opts.AbortTimeout c.network = opts.Network c.remotePeerID = opts.RemotePeerID + + if opts.BaseContext != nil { + c.baseContext = opts.BaseContext + } } if c.abortTimeout == 0 { c.abortTimeout = 100 * time.Millisecond @@ -261,7 +272,7 @@ func NewConn(t Transport, opts *Options) *Conn { func (c *Conn) startBackgroundTasks() { // We use an errgroup to link the lifetime of background tasks // to each other. - ctx, cancel := context.WithCancel(context.Background()) + ctx, cancel := context.WithCancel(c.baseContext()) g, ctx := errgroup.WithContext(ctx) c.bgctx = ctx From 376c32ad33400786a9c69889340e26d39f3e5d2b Mon Sep 17 00:00:00 2001 From: Dzmitry Mikhalapau Date: Thu, 28 Nov 2024 21:23:52 +0500 Subject: [PATCH 2/2] Fix typing errors Signed-off-by: Dzmitry Mikhalapau --- rpc/rpc.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/rpc/rpc.go b/rpc/rpc.go index fedd76d6..c9ccdb46 100644 --- a/rpc/rpc.go +++ b/rpc/rpc.go @@ -91,7 +91,7 @@ type Conn struct { baseContext func() context.Context // bgctx is a Context that is canceled when shutdown starts. Note - // that if baseContext is not provided. it's parent is context.Background(), + // that if baseContext is not provided, it's parent is context.Background(), // so we can rely on this being the *only* time it will be canceled. bgctx context.Context @@ -204,7 +204,7 @@ type Options struct { // set this. Network Network - // BaseContext is an optional funcation that returns a base context + // BaseContext is an optional function that returns a base context // for any incoming connection. If ommitted, the context.Background() // will be used instead. BaseContext func() context.Context