From 26c1c3f2883b74711c46775b250ecbdd8093a719 Mon Sep 17 00:00:00 2001
From: Vivek Menezes <vivek@cockroachlabs.com>
Date: Wed, 22 Aug 2018 08:46:36 -0400
Subject: [PATCH] sql: fix flaky TestLeaseRenewedPeriodically

I changed GossipUpdateEvent as well which is strictly not needed
but there is a chance without it there will continue to be
some flakiness.

This fix essentially moves the check for release counts
further up before lease acquisition. This fixes a race
where after lease acquisition the periodic timer would
fire, acquire/release some leases.

fixes #28771

Release note: None
---
 pkg/sql/lease.go               |  9 ++++++---
 pkg/sql/lease_internal_test.go |  3 ++-
 pkg/sql/lease_test.go          | 20 ++++++++++++--------
 3 files changed, 20 insertions(+), 12 deletions(-)

diff --git a/pkg/sql/lease.go b/pkg/sql/lease.go
index 7933a85eb542..8a37e4f6e56f 100644
--- a/pkg/sql/lease.go
+++ b/pkg/sql/lease.go
@@ -1168,8 +1168,9 @@ var _ base.ModuleTestingKnobs = &LeaseStoreTestingKnobs{}
 type LeaseManagerTestingKnobs struct {
 	// A callback called when a gossip update is received, before the leases are
 	// refreshed. Careful when using this to block for too long - you can block
-	// all the gossip users in the system.
-	GossipUpdateEvent func(config.SystemConfig)
+	// all the gossip users in the system. If it returns an error the gossip
+	// update is ignored.
+	GossipUpdateEvent func(config.SystemConfig) error
 	// A callback called after the leases are refreshed as a result of a gossip update.
 	TestingLeasesRefreshedEvent func(config.SystemConfig)
 
@@ -1627,7 +1628,9 @@ func (m *LeaseManager) RefreshLeases(s *stop.Stopper, db *client.DB, g *gossip.G
 			case <-gossipUpdateC:
 				cfg, _ := g.GetSystemConfig()
 				if m.testingKnobs.GossipUpdateEvent != nil {
-					m.testingKnobs.GossipUpdateEvent(cfg)
+					if err := m.testingKnobs.GossipUpdateEvent(cfg); err != nil {
+						break
+					}
 				}
 				// Read all tables and their versions
 				if log.V(2) {
diff --git a/pkg/sql/lease_internal_test.go b/pkg/sql/lease_internal_test.go
index 6fb940b68b5c..a2bef11137f7 100644
--- a/pkg/sql/lease_internal_test.go
+++ b/pkg/sql/lease_internal_test.go
@@ -127,9 +127,10 @@ func TestPurgeOldVersions(t *testing.T) {
 	serverParams := base.TestServerArgs{
 		Knobs: base.TestingKnobs{
 			SQLLeaseManager: &LeaseManagerTestingKnobs{
-				GossipUpdateEvent: func(cfg config.SystemConfig) {
+				GossipUpdateEvent: func(cfg config.SystemConfig) error {
 					gossipSem <- struct{}{}
 					<-gossipSem
+					return nil
 				},
 			},
 		},
diff --git a/pkg/sql/lease_test.go b/pkg/sql/lease_test.go
index 0c55e674b380..bed1946806e5 100644
--- a/pkg/sql/lease_test.go
+++ b/pkg/sql/lease_test.go
@@ -1591,6 +1591,9 @@ func TestLeaseRenewedPeriodically(testingT *testing.T) {
 					atomic.AddInt32(&testAcquisitionBlockCount, 1)
 				},
 			},
+			GossipUpdateEvent: func(cfg config.SystemConfig) error {
+				return errors.Errorf("ignore gossip update")
+			},
 		},
 	}
 	params.LeaseManagerConfig = base.NewLeaseManagerConfig()
@@ -1620,6 +1623,15 @@ CREATE TABLE t.test2 ();
 
 	atomic.StoreInt32(&testAcquisitionBlockCount, 0)
 
+	numReleasedLeases := func() int {
+		mu.Lock()
+		defer mu.Unlock()
+		return len(releasedIDs)
+	}
+	if count := numReleasedLeases(); count != 0 {
+		t.Fatalf("expected no leases to be releases, released %d", count)
+	}
+
 	// Acquire a lease on test1 by name.
 	ts1, _, err := t.node(1).AcquireByName(ctx, t.server.Clock().Now(), dbID, "test1")
 	if err != nil {
@@ -1644,14 +1656,6 @@ CREATE TABLE t.test2 ();
 
 	// From now on henceforth do not acquire a lease, so any renewals can only
 	// happen through the periodic lease renewal mechanism.
-	numReleasedLeases := func() int {
-		mu.Lock()
-		defer mu.Unlock()
-		return len(releasedIDs)
-	}
-	if count := numReleasedLeases(); count != 0 {
-		t.Fatalf("expected no leases to be releases, released %d", count)
-	}
 
 	// Reset testAcquisitionBlockCount as the first acqusitions will always block.
 	atomic.StoreInt32(&testAcquisitionBlockCount, 0)