Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Provide an option to specify base context for RPC connection #595

Open
wants to merge 2 commits into
base: main
Choose a base branch
from
Open

Provide an option to specify base context for RPC connection #595

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
c67d45b
Provide an option to specify base context for RPC connection
homier Nov 27, 2024
376c32a
Fix typing errors
homier Nov 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 16 additions & 5 deletions rpc/rpc.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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 function 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
Expand Down Expand Up @@ -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]()
Expand All @@ -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
Expand All @@ -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
Expand Down