From 2331ca1d184540c0108b3ce6544b6f6d3a5b7204 Mon Sep 17 00:00:00 2001
From: Joey Pereira <joey@pereira.io>
Date: Tue, 10 Oct 2017 00:05:11 -0400
Subject: [PATCH] sql: Refactor LeaseDuration to prevent test collisions.

Previously, the global LeaseDuration value was modified during a test.
If this test is run in parallel, this would cause problems. The global
has now been modified to live on the LeaseStore where it can be modified
in tests safely for the life of the test.
---
 pkg/sql/lease.go      | 30 +++++++++++++++++++++---------
 pkg/sql/lease_test.go | 26 +++++++++++++-------------
 2 files changed, 34 insertions(+), 22 deletions(-)

diff --git a/pkg/sql/lease.go b/pkg/sql/lease.go
index ab78d9999559..25924c2ccfe6 100644
--- a/pkg/sql/lease.go
+++ b/pkg/sql/lease.go
@@ -44,7 +44,7 @@ import (
 // TODO(pmattis): Periodically renew leases for tables that were used recently and
 // for which the lease will expire soon.
 
-var (
+const (
 	// LeaseDuration is the mean duration a lease will be acquired for. The
 	// actual duration is jittered in the range
 	// [0.75,1.25]*LeaseDuration. Exported for testing purposes only.
@@ -118,14 +118,18 @@ type LeaseStore struct {
 	clock  *hlc.Clock
 	nodeID *base.NodeIDContainer
 
+	// leaseDuration is a constant initialized by NewLeaseManager. It is modified
+	// only during by tests through testSetLeaseDuration.
+	leaseDuration time.Duration
+
 	testingKnobs LeaseStoreTestingKnobs
 	memMetrics   *MemoryMetrics
 }
 
 // jitteredLeaseDuration returns a randomly jittered duration from the interval
 // [0.75 * leaseDuration, 1.25 * leaseDuration].
-func jitteredLeaseDuration() time.Duration {
-	return time.Duration(float64(LeaseDuration) * (0.75 + 0.5*rand.Float64()))
+func jitteredLeaseDuration(leaseDuration time.Duration) time.Duration {
+	return time.Duration(float64(leaseDuration) * (0.75 + 0.5*rand.Float64()))
 }
 
 // acquire a lease on the most recent version of a table descriptor.
@@ -137,7 +141,7 @@ func (s LeaseStore) acquire(
 	var table *tableVersionState
 	err := s.db.Txn(ctx, func(ctx context.Context, txn *client.Txn) error {
 		expiration := txn.OrigTimestamp()
-		expiration.WallTime += int64(jitteredLeaseDuration())
+		expiration.WallTime += int64(jitteredLeaseDuration(s.leaseDuration))
 		if expiration.Less(minExpirationTime) {
 			expiration = minExpirationTime
 		}
@@ -1099,11 +1103,12 @@ func NewLeaseManager(
 ) *LeaseManager {
 	lm := &LeaseManager{
 		LeaseStore: LeaseStore{
-			db:           db,
-			clock:        clock,
-			nodeID:       nodeID,
-			testingKnobs: testingKnobs.LeaseStoreTestingKnobs,
-			memMetrics:   memMetrics,
+			db:            db,
+			clock:         clock,
+			nodeID:        nodeID,
+			leaseDuration: LeaseDuration,
+			testingKnobs:  testingKnobs.LeaseStoreTestingKnobs,
+			memMetrics:    memMetrics,
 		},
 		testingKnobs: testingKnobs,
 		tableNames: tableNameCache{
@@ -1124,6 +1129,13 @@ func nameMatchesTable(table *sqlbase.TableDescriptor, dbID sqlbase.ID, tableName
 	return table.ParentID == dbID && table.Name == tableName
 }
 
+// TestSetLeaseDuration is used to modify the leaseDuration constant for
+// the lifespan of the LeaseManager. This function exists and is exported only
+// for testing purposes.
+func (m *LeaseManager) TestSetLeaseDuration(leaseDuration time.Duration) {
+	m.LeaseStore.leaseDuration = leaseDuration
+}
+
 // AcquireByName returns a table version for the specified table valid for
 // the timestamp. It returns the table descriptor and a expiration time.
 // A transaction using this descriptor must ensure that its
diff --git a/pkg/sql/lease_test.go b/pkg/sql/lease_test.go
index e5c79a1c149a..f8281169b731 100644
--- a/pkg/sql/lease_test.go
+++ b/pkg/sql/lease_test.go
@@ -50,16 +50,18 @@ type leaseTest struct {
 	kvDB                     *client.DB
 	nodes                    map[uint32]*sql.LeaseManager
 	leaseManagerTestingKnobs sql.LeaseManagerTestingKnobs
+	leaseDuration            time.Duration
 }
 
 func newLeaseTest(t *testing.T, params base.TestServerArgs) *leaseTest {
 	s, db, kvDB := serverutils.StartServer(t, params)
 	leaseTest := &leaseTest{
-		T:      t,
-		server: s,
-		db:     db,
-		kvDB:   kvDB,
-		nodes:  map[uint32]*sql.LeaseManager{},
+		T:             t,
+		server:        s,
+		db:            db,
+		kvDB:          kvDB,
+		nodes:         map[uint32]*sql.LeaseManager{},
+		leaseDuration: sql.LeaseDuration,
 	}
 	if params.Knobs.SQLLeaseManager != nil {
 		leaseTest.leaseManagerTestingKnobs =
@@ -109,6 +111,10 @@ func (t *leaseTest) expectLeases(descID sqlbase.ID, expected string) {
 	})
 }
 
+func (t *leaseTest) setLeaseDuration(leaseDuration time.Duration) {
+	t.leaseDuration = leaseDuration
+}
+
 func (t *leaseTest) acquire(
 	nodeID uint32, descID sqlbase.ID,
 ) (*sqlbase.TableDescriptor, hlc.Timestamp, error) {
@@ -192,6 +198,7 @@ func (t *leaseTest) node(nodeID uint32) *sql.LeaseManager {
 			t.server.Stopper(),
 			&sql.MemoryMetrics{},
 		)
+		mgr.TestSetLeaseDuration(t.leaseDuration)
 		t.nodes[nodeID] = mgr
 	}
 	return mgr
@@ -318,14 +325,7 @@ func TestLeaseManagerReacquire(testingT *testing.T) {
 
 	// Set the lease duration such that the next lease acquisition will
 	// require the lease to be reacquired.
-	savedLeaseDuration := sql.LeaseDuration
-	defer func() {
-		sql.LeaseDuration = savedLeaseDuration
-	}()
-
-	sql.LeaseDuration = 5 * time.Nanosecond
-
-	time.Sleep(5 * sql.LeaseDuration)
+	t.setLeaseDuration(5 * time.Nanosecond)
 
 	l1, e1 := t.mustAcquire(1, descID)
 	t.expectLeases(descID, "/1/1")