From 1e330a90c77d14fbf3af7e4a47e6ab064eedd0c9 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Thu, 20 Oct 2016 17:03:53 -0700
Subject: [PATCH 1/2] concurrency: terminate session.Close if revoke takes
 longer than TTL

Fixes #6681
---
 clientv3/concurrency/session.go | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/clientv3/concurrency/session.go b/clientv3/concurrency/session.go
index fb4fd1162..5f0c7776e 100644
--- a/clientv3/concurrency/session.go
+++ b/clientv3/concurrency/session.go
@@ -15,6 +15,8 @@
 package concurrency
 
 import (
+	"time"
+
 	v3 "github.com/coreos/etcd/clientv3"
 	"golang.org/x/net/context"
 )
@@ -25,6 +27,7 @@ const defaultSessionTTL = 60
 // Fault-tolerant applications may use sessions to reason about liveness.
 type Session struct {
 	client *v3.Client
+	opts   *sessionOptions
 	id     v3.LeaseID
 
 	cancel context.CancelFunc
@@ -51,7 +54,7 @@ func NewSession(client *v3.Client, opts ...SessionOption) (*Session, error) {
 	}
 
 	donec := make(chan struct{})
-	s := &Session{client: client, id: id, cancel: cancel, donec: donec}
+	s := &Session{client: client, opts: ops, id: id, cancel: cancel, donec: donec}
 
 	// keep the lease alive until client error or cancelled context
 	go func() {
@@ -87,7 +90,10 @@ func (s *Session) Orphan() {
 // Close orphans the session and revokes the session lease.
 func (s *Session) Close() error {
 	s.Orphan()
-	_, err := s.client.Revoke(s.client.Ctx(), s.id)
+	// 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)
+	_, err := s.client.Revoke(ctx, s.id)
+	cancel()
 	return err
 }
 

From f38a5d19a8652c5db9f7f48a1bacbdd38a631139 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Fri, 21 Oct 2016 15:12:11 -0700
Subject: [PATCH 2/2] concurrency: add WithContext option to sessions

Makes it possible to cancel session requests without having to
close the entire client.
---
 clientv3/concurrency/session.go | 20 ++++++++++++++++----
 1 file changed, 16 insertions(+), 4 deletions(-)

diff --git a/clientv3/concurrency/session.go b/clientv3/concurrency/session.go
index 5f0c7776e..59267df50 100644
--- a/clientv3/concurrency/session.go
+++ b/clientv3/concurrency/session.go
@@ -36,18 +36,18 @@ type Session struct {
 
 // NewSession gets the leased session for a client.
 func NewSession(client *v3.Client, opts ...SessionOption) (*Session, error) {
-	ops := &sessionOptions{ttl: defaultSessionTTL}
+	ops := &sessionOptions{ttl: defaultSessionTTL, ctx: client.Ctx()}
 	for _, opt := range opts {
 		opt(ops)
 	}
 
-	resp, err := client.Grant(client.Ctx(), int64(ops.ttl))
+	resp, err := client.Grant(ops.ctx, int64(ops.ttl))
 	if err != nil {
 		return nil, err
 	}
 	id := v3.LeaseID(resp.ID)
 
-	ctx, cancel := context.WithCancel(client.Ctx())
+	ctx, cancel := context.WithCancel(ops.ctx)
 	keepAlive, err := client.KeepAlive(ctx, id)
 	if err != nil || keepAlive == nil {
 		return nil, err
@@ -91,7 +91,7 @@ func (s *Session) Orphan() {
 func (s *Session) Close() error {
 	s.Orphan()
 	// 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)
 	cancel()
 	return err
@@ -99,6 +99,7 @@ func (s *Session) Close() error {
 
 type sessionOptions struct {
 	ttl int
+	ctx context.Context
 }
 
 // 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
+	}
+}