Repository: snowflakedb/gosnowflake
Pull request #398 (2021-03-09 08:05:53)

SNOW-291628 Synchronize the session renewal when token expires
### Description
When a request comes back as session expired, we issue a
session renewal and rotate the token in the token accessor.

Multiple requests can hit the same session expired error at
the same time, and trigger multiple renewals of the same stale token.

Use a mutex to synchronize the token renewal, and only issue
one renewal for the stale token.

Testing

New tests with multiple goroutines added to verify concurrency.



### Checklist
- [x] Code compiles correctly
- [x] Run ``make fmt`` to fix inconsistent formats
- [x] Run ``make lint`` to get lint errors and fix all of them
- [x] Created tests which fail without the change (if possible)
- [x] All tests passing
- [ ] Extended the README / documentation, if necessary


sfc-gh-dwu: Going to move the mutex lock & unlock into the token accessor interface so that the driver user can control synchronizing across servers for renewals.
sfc-gh-anezvigin: Left a couple thoughts, otherwise LGTM.
sfc-gh-dwu: Added new tests to verify:
- with no contention, renewals can happen
- with controlled locking, if token is swapped while renewal goroutine is locked, it will check and not renew because it became stale
- contention with many goroutines trying to compete for renewal, only one wins

Patch:
diff --git a/heartbeat.go b/heartbeat.go
index 3e7b203..fb39ef7 100644
--- a/heartbeat.go
+++ b/heartbeat.go
@@ -78,7 +78,7 @@ func (hc *heartbeat) heartbeatMain() error {
 			return err
 		}
 		if respd.Code == sessionExpiredCode {
-			err = hc.restful.FuncRenewSession(context.TODO(), hc.restful, timeout)
+			err = hc.restful.renewExpiredSessionToken(context.Background(), timeout, token)
 			if err != nil {
 				return err
 			}
diff --git a/restful.go b/restful.go
index 36a8a1b..7cee0e9 100644
--- a/restful.go
+++ b/restful.go
@@ -87,6 +87,21 @@ func (sr *snowflakeRestful) getFullURL(path string, params *url.Values) *url.URL
 	return ret
 }
 
+// Renew the snowflake session if the current token is still the stale token specified
+func (sr *snowflakeRestful) renewExpiredSessionToken(ctx context.Context, timeout time.Duration, expiredToken string) error {
+	err := sr.TokenAccessor.Lock()
+	if err != nil {
+		return err
+	}
+	defer sr.TokenAccessor.Unlock()
+	currentToken, _, _ := sr.TokenAccessor.GetTokens()
+	if expiredToken == currentToken || currentToken == "" {
+		// Only renew the session if the current token is still the expired token or current token is empty
+		return sr.FuncRenewSession(ctx, sr, timeout)
+	}
+	return nil
+}
+
 type renewSessionResponse struct {
 	Data    renewSessionResponseMain `json:"data"`
 	Message string                   `json:"message"`
@@ -195,7 +210,7 @@ func postRestfulQueryHelper(
 			return nil, err
 		}
 		if respd.Code == sessionExpiredCode {
-			err = sr.FuncRenewSession(ctx, sr, timeout)
+			err = sr.renewExpiredSessionToken(ctx, timeout, token)
 			if err != nil {
 				return nil, err
 			}
@@ -260,7 +275,7 @@ func postRestfulQueryHelper(
 				return nil, err
 			}
 			if respd.Code == sessionExpiredCode {
-				err = sr.FuncRenewSession(ctx, sr, timeout)
+				err = sr.renewExpiredSessionToken(ctx, timeout, token)
 				if err != nil {
 					return nil, err
 				}
diff --git a/util.go b/util.go
index a568d61..6e1e118 100644
--- a/util.go
+++ b/util.go
@@ -6,6 +6,7 @@ import (
 	"context"
 	"database/sql/driver"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/google/uuid"
@@ -118,23 +119,40 @@ func toNamedValues(values []driver.Value) []driver.NamedValue {
 type TokenAccessor interface {
 	GetTokens() (token string, masterToken string, sessionID int)
 	SetTokens(token string, masterToken string, sessionID int)
+	Lock() error
+	Unlock()
 }
 
 type simpleTokenAccessor struct {
-	token       string
-	masterToken string
-	sessionID   int
+	token        string
+	masterToken  string
+	sessionID    int
+	accessorLock sync.Mutex   // Used to implement accessor's Lock and Unlock
+	tokenLock    sync.RWMutex // Used to synchronize SetTokens and GetTokens
 }
 
 func getSimpleTokenAccessor() TokenAccessor {
 	return &simpleTokenAccessor{sessionID: -1}
 }
 
+func (sta *simpleTokenAccessor) Lock() error {
+	sta.accessorLock.Lock()
+	return nil
+}
+
+func (sta *simpleTokenAccessor) Unlock() {
+	sta.accessorLock.Unlock()
+}
+
 func (sta *simpleTokenAccessor) GetTokens() (token string, masterToken string, sessionID int) {
+	sta.tokenLock.RLock()
+	defer sta.tokenLock.RUnlock()
 	return sta.token, sta.masterToken, sta.sessionID
 }
 
 func (sta *simpleTokenAccessor) SetTokens(token string, masterToken string, sessionID int) {
+	sta.tokenLock.Lock()
+	defer sta.tokenLock.Unlock()
 	sta.token = token
 	sta.masterToken = masterToken
 	sta.sessionID = sessionID


Test patch:
diff --git a/restful_test.go b/restful_test.go
index cbb4ce3..d6dcb6a 100644
--- a/restful_test.go
+++ b/restful_test.go
@@ -9,6 +9,8 @@ import (
 	"fmt"
 	"net/http"
 	"net/url"
+	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -124,6 +126,190 @@ func renewSessionTestError(_ context.Context, _ *snowflakeRestful, _ time.Durati
 	return errors.New("failed to renew session in tests")
 }
 
+func TestUnitTokenAccessorDoesNotRenewStaleToken(t *testing.T) {
+	accessor := getSimpleTokenAccessor()
+	oldToken := "test"
+	accessor.SetTokens(oldToken, "master", 123)
+
+	renewSessionCalled := false
+	renewSessionDummy := func(_ context.Context, sr *snowflakeRestful, _ time.Duration) error {
+		// should not have gotten to actual renewal
+		renewSessionCalled = true
+		return nil
+	}
+
+	sr := &snowflakeRestful{
+		FuncRenewSession: renewSessionDummy,
+		TokenAccessor:    accessor,
+	}
+
+	// try to intentionally renew with stale token
+	sr.renewExpiredSessionToken(context.Background(), time.Hour, "stale-token")
+
+	if renewSessionCalled {
+		t.Fatal("FuncRenewSession should not have been called")
+	}
+
+	// set the current token to empty, should still call renew even if stale token is passed in
+	accessor.SetTokens("", "master", 123)
+	sr.renewExpiredSessionToken(context.Background(), time.Hour, "stale-token")
+
+	if !renewSessionCalled {
+		t.Fatal("FuncRenewSession should have been called because current token is empty")
+	}
+}
+
+type wrappedAccessor struct {
+	ta              TokenAccessor
+	lockCallCount   int32
+	unlockCallCount int32
+}
+
+func (wa *wrappedAccessor) Lock() error {
+	atomic.AddInt32(&wa.lockCallCount, 1)
+	err := wa.ta.Lock()
+	return err
+}
+
+func (wa *wrappedAccessor) Unlock() {
+	atomic.AddInt32(&wa.unlockCallCount, 1)
+	wa.ta.Unlock()
+}
+
+func (wa *wrappedAccessor) GetTokens() (token string, masterToken string, sessionID int) {
+	return wa.ta.GetTokens()
+}
+
+func (wa *wrappedAccessor) SetTokens(token string, masterToken string, sessionID int) {
+	wa.ta.SetTokens(token, masterToken, sessionID)
+}
+
+func TestUnitTokenAccessorRenewBlocked(t *testing.T) {
+	accessor := wrappedAccessor{
+		ta: getSimpleTokenAccessor(),
+	}
+	oldToken := "test"
+	accessor.SetTokens(oldToken, "master", 123)
+
+	renewSessionCalled := false
+	renewSessionDummy := func(_ context.Context, sr *snowflakeRestful, _ time.Duration) error {
+		renewSessionCalled = true
+		return nil
+	}
+
+	sr := &snowflakeRestful{
+		FuncRenewSession: renewSessionDummy,
+		TokenAccessor:    &accessor,
+	}
+
+	// intentionally lock the accessor first
+	accessor.Lock()
+
+	// try to intentionally renew with stale token
+	var renewalStart sync.WaitGroup
+	var renewalDone sync.WaitGroup
+	renewalStart.Add(1)
+	renewalDone.Add(1)
+	go func() {
+		renewalStart.Done()
+		sr.renewExpiredSessionToken(context.Background(), time.Hour, oldToken)
+		renewalDone.Done()
+	}()
+
+	// wait for renewal to start and get blocked on lock
+	renewalStart.Wait()
+	// should be blocked and not be able to call renew session
+	if renewSessionCalled {
+		t.Fail()
+	}
+
+	// rotate the token again so that the session token is considered stale
+	accessor.SetTokens("new-token", "m", 321)
+
+	// unlock so that renew can happen
+	accessor.Unlock()
+	renewalDone.Wait()
+
+	// renewal should be done but token should still not
+	// have been renewed since we intentionally swapped token while locked
+	if renewSessionCalled {
+		t.Fail()
+	}
+
+	// wait for accessor defer unlock
+	accessor.Lock()
+	if accessor.lockCallCount != 3 {
+		t.Fatalf("Expected Lock() to be called thrice, but got %v", accessor.lockCallCount)
+	}
+	if accessor.unlockCallCount != 2 {
+		t.Fatalf("Expected Unlock() to be called twice, but got %v", accessor.unlockCallCount)
+	}
+}
+
+func TestUnitTokenAccessorRenewSessionContention(t *testing.T) {
+	accessor := getSimpleTokenAccessor()
+	oldToken := "test"
+	accessor.SetTokens(oldToken, "master", 123)
+	var counter int32 = 0
+
+	expectedToken := "new token"
+	expectedMaster := "new master"
+	expectedSession := 321
+
+	renewSessionDummy := func(_ context.Context, sr *snowflakeRestful, _ time.Duration) error {
+		accessor.SetTokens(expectedToken, expectedMaster, expectedSession)
+		atomic.AddInt32(&counter, 1)
+		return nil
+	}
+
+	sr := &snowflakeRestful{
+		FuncRenewSession: renewSessionDummy,
+		TokenAccessor:    accessor,
+	}
+
+	var renewalsStart sync.WaitGroup
+	var renewalsDone sync.WaitGroup
+	var renewalError error
+	numRoutines := 50
+	for i := 0; i < numRoutines; i++ {
+		renewalsDone.Add(1)
+		renewalsStart.Add(1)
+		go func() {
+			// wait for all goroutines to have been created before proceeding to race against each other
+			renewalsStart.Wait()
+			err := sr.renewExpiredSessionToken(context.Background(), time.Hour, oldToken)
+			if err != nil {
+				renewalError = err
+			}
+			renewalsDone.Done()
+		}()
+	}
+
+	// unlock all of the waiting goroutines simultaneously
+	renewalsStart.Add(-numRoutines)
+
+	// wait for all competing goroutines to finish calling renew expired session token
+	renewalsDone.Wait()
+
+	if renewalError != nil {
+		t.Fatalf("failed to renew session, error %v", renewalError)
+	}
+	newToken, newMaster, newSession := accessor.GetTokens()
+	if newToken != expectedToken {
+		t.Fatalf("token %v does not match expected %v", newToken, expectedToken)
+	}
+	if newMaster != expectedMaster {
+		t.Fatalf("master token %v does not match expected %v", newMaster, expectedMaster)
+	}
+	if newSession != expectedSession {
+		t.Fatalf("session %v does not match expected %v", newSession, expectedSession)
+	}
+	// only the first renewal will go through and FuncRenewSession should be called exactly once
+	if counter != 1 {
+		t.Fatalf("renew expired session was called more than once: %v", counter)
+	}
+}
+
 func TestUnitPostQueryHelperUsesToken(t *testing.T) {
 	accessor := getSimpleTokenAccessor()
 	token := "token123"
diff --git a/util_test.go b/util_test.go
index ebc8c3c..1e78165 100644
--- a/util_test.go
+++ b/util_test.go
@@ -5,6 +5,9 @@ package gosnowflake
 import (
 	"context"
 	"database/sql/driver"
+	"math/rand"
+	"strconv"
+	"sync"
 	"testing"
 	"time"
 
@@ -44,6 +47,34 @@ func TestSimpleTokenAccessor(t *testing.T) {
 	}
 }
 
+func TestSimpleTokenAccessorGetTokensSynchronization(t *testing.T) {
+	accessor := getSimpleTokenAccessor()
+	var wg sync.WaitGroup
+	failed := false
+	for i := 0; i < 1000; i++ {
+		wg.Add(1)
+		go func() {
+			// set a random session and token
+			session := rand.Int()
+			sessionStr := strconv.Itoa(session)
+			accessor.SetTokens("t"+sessionStr, "m"+sessionStr, session)
+
+			// read back session and token and verify that invariant still holds
+			token, masterToken, session := accessor.GetTokens()
+			sessionStr = strconv.Itoa(session)
+			if "t"+sessionStr != token || "m"+sessionStr != masterToken {
+				failed = true
+			}
+			wg.Done()
+		}()
+	}
+	// wait for all competing goroutines to finish setting and getting tokens
+	wg.Wait()
+	if failed {
+		t.Fail()
+	}
+}
+
 func TestGetRequestIDFromContext(t *testing.T) {
 	expectedRequestID := uuid.New()
 	ctx := WithRequestID(context.Background(), expectedRequestID)