concurrency: add WithContext option to sessions

Makes it possible to cancel session requests without having to
close the entire client.
This commit is contained in:
Anthony Romano 2016-10-21 15:12:11 -07:00
parent 1e330a90c7
commit f38a5d19a8

View File

@ -36,18 +36,18 @@ type Session struct {
// NewSession gets the leased session for a client. // NewSession gets the leased session for a client.
func NewSession(client *v3.Client, opts ...SessionOption) (*Session, error) { func NewSession(client *v3.Client, opts ...SessionOption) (*Session, error) {
ops := &sessionOptions{ttl: defaultSessionTTL} ops := &sessionOptions{ttl: defaultSessionTTL, ctx: client.Ctx()}
for _, opt := range opts { for _, opt := range opts {
opt(ops) opt(ops)
} }
resp, err := client.Grant(client.Ctx(), int64(ops.ttl)) resp, err := client.Grant(ops.ctx, int64(ops.ttl))
if err != nil { if err != nil {
return nil, err return nil, err
} }
id := v3.LeaseID(resp.ID) id := v3.LeaseID(resp.ID)
ctx, cancel := context.WithCancel(client.Ctx()) ctx, cancel := context.WithCancel(ops.ctx)
keepAlive, err := client.KeepAlive(ctx, id) keepAlive, err := client.KeepAlive(ctx, id)
if err != nil || keepAlive == nil { if err != nil || keepAlive == nil {
return nil, err return nil, err
@ -91,7 +91,7 @@ func (s *Session) Orphan() {
func (s *Session) Close() error { func (s *Session) Close() error {
s.Orphan() s.Orphan()
// if revoke takes longer than the ttl, lease is expired anyway // if revoke takes longer than the ttl, lease is expired anyway
ctx, cancel := context.WithTimeout(s.client.Ctx(), time.Duration(s.opts.ttl)*time.Second) ctx, cancel := context.WithTimeout(s.opts.ctx, time.Duration(s.opts.ttl)*time.Second)
_, err := s.client.Revoke(ctx, s.id) _, err := s.client.Revoke(ctx, s.id)
cancel() cancel()
return err return err
@ -99,6 +99,7 @@ func (s *Session) Close() error {
type sessionOptions struct { type sessionOptions struct {
ttl int ttl int
ctx context.Context
} }
// SessionOption configures Session. // SessionOption configures Session.
@ -113,3 +114,14 @@ func WithTTL(ttl int) SessionOption {
} }
} }
} }
// WithContext assigns a context to the session instead of defaulting to
// using the client context. This is useful for canceling NewSession and
// Close operations immediately without having to close the client. If the
// context is canceled before Close() completes, the session's lease will be
// abandoned and left to expire instead of being revoked.
func WithContext(ctx context.Context) SessionOption {
return func(so *sessionOptions) {
so.ctx = ctx
}
}