gossip: pass OrigStamp as part of gossip callback

Previously, gossip callbacks were registered with a key string to receive
updates when new gossip info for that key arrived. These callbacks took the key
and value as arguments. This commit adds the originating timestamp (from
gossip.Info) as an additional argument, representing the wall time when the info
was first generated by the originating node. The new argument is not used by
existing code, but mma will use it in future commits to determine whether store
load info reflects certain replication changes.

Epic: none
Release note: none

Author: wenyihu6

pkg/gossip/client.go

@@ -312,7 +312,7 @@ func (c *client) gossip(
 	sendGossipChan := make(chan struct{}, 1)
 
 	// Register a callback for gossip updates.
-	updateCallback := func(_ string, _ roachpb.Value) {
+	updateCallback := func(_ string, _ roachpb.Value, _ int64) {
 		select {
 		case sendGossipChan <- struct{}{}:
 		default:

pkg/gossip/gossip.go

@@ -756,7 +756,7 @@ func maxPeers(nodeCount int) int {
 // new addresses for each encountered host and to write the
 // set of gossip node addresses to persistent storage when it
 // changes.
-func (g *Gossip) updateNodeAddress(key string, content roachpb.Value) {
+func (g *Gossip) updateNodeAddress(key string, content roachpb.Value, _ int64) {
 	ctx := g.AnnotateCtx(context.TODO())
 	var desc roachpb.NodeDescriptor
 	if err := content.GetProto(&desc); err != nil {
@@ -825,7 +825,7 @@ func (g *Gossip) removeNodeDescriptorLocked(nodeID roachpb.NodeID) {
 }
 
 // updateStoreMap is a gossip callback which is used to update storeMap.
-func (g *Gossip) updateStoreMap(key string, content roachpb.Value) {
+func (g *Gossip) updateStoreMap(key string, content roachpb.Value, _ int64) {
 	ctx := g.AnnotateCtx(context.TODO())
 	var desc roachpb.StoreDescriptor
 	if err := content.GetProto(&desc); err != nil {
@@ -1133,9 +1133,10 @@ func (g *Gossip) IterateInfos(prefix string, visit func(k string, info Info) err
 	return nil
 }
 
-// Callback is a callback method to be invoked on gossip update
-// of info denoted by key.
-type Callback func(string, roachpb.Value)
+// Callback is a callback method to be invoked on gossip update of info denoted
+// by key. origTimestamp is the info wall time when generated by the originating
+// node.
+type Callback func(key string, value roachpb.Value, origTimestamp int64)
 
 // CallbackOption is a marker interface that callback options must implement.
 type CallbackOption interface {
@@ -1616,7 +1617,7 @@ func (g *Gossip) TestingAddInfoProtoAndWaitForAllCallbacks(
 		if cb.matcher.MatchString(key) {
 			cb.cw.mu.Lock()
 			cb.cw.mu.workQueue = append(cb.cw.mu.workQueue, callbackWorkItem{
-				method: func(_ string, _ roachpb.Value) {
+				method: func(_ string, _ roachpb.Value, _ int64) {
 					wg.Done()
 				},
 				schedulingTime: timeutil.Now(),
@@ -1657,7 +1658,7 @@ func (g *Gossip) GetFirstRangeDescriptor() (*roachpb.RangeDescriptor, error) {
 
 // OnFirstRangeChanged implements kvcoord.FirstRangeProvider.
 func (g *Gossip) OnFirstRangeChanged(cb func(*roachpb.RangeDescriptor)) {
-	g.RegisterCallback(KeyFirstRangeDescriptor, func(_ string, value roachpb.Value) {
+	g.RegisterCallback(KeyFirstRangeDescriptor, func(_ string, value roachpb.Value, _ int64) {
 		ctx := context.Background()
 		desc := &roachpb.RangeDescriptor{}
 		if err := value.GetProto(desc); err != nil {

pkg/gossip/gossip_test.go

@@ -1172,7 +1172,7 @@ func TestCallbacksPendingMetricGoesToZeroOnStop(t *testing.T) {
 			defer stopper.Stop(ctx)
 			g := NewTest(1, stopper, metric.NewRegistry())
 
-			unregister := g.RegisterCallback("test.*", func(key string, val roachpb.Value) {
+			unregister := g.RegisterCallback("test.*", func(key string, val roachpb.Value, _ int64) {
 				// Do nothing.
 			})
 

pkg/gossip/infostore.go

@@ -79,8 +79,11 @@ type callbackWorkItem struct {
 	// schedulingTime is the time when the callback was scheduled.
 	schedulingTime time.Time
 	method         Callback
-	key            string
-	content        roachpb.Value
+	// key, content, origTimestamp are the parameters that will be passed to the
+	// callback method. They are based on the infos added to the infostore.
+	key           string
+	content       roachpb.Value
+	origTimestamp int64
 }
 
 type callbackWork struct {
@@ -274,7 +277,7 @@ func (is *infoStore) launchCallbackWorker(ambient log.AmbientContext, cw *callba
 						is.metrics.CallbacksPendingDuration.RecordValue(queueDur.Nanoseconds())
 					}
 
-					work.method(work.key, work.content)
+					work.method(work.key, work.content, work.origTimestamp)
 
 					afterProcess := timeutil.Now()
 					processDur := afterProcess.Sub(afterQueue)
@@ -362,7 +365,7 @@ func (is *infoStore) addInfo(key string, i *Info) error {
 	ratchetHighWaterStamp(is.highWaterStamps, i.NodeID, i.OrigStamp)
 	changed := existingInfo == nil ||
 		!bytes.Equal(existingInfo.Value.RawBytes, i.Value.RawBytes)
-	is.processCallbacks(key, i.Value, changed)
+	is.processCallbacks(key, i.Value, i.OrigStamp, changed)
 	return nil
 }
 
@@ -419,7 +422,7 @@ func (is *infoStore) registerCallback(
 
 	if err := is.visitInfos(func(key string, i *Info) error {
 		if matcher.MatchString(key) {
-			is.runCallbacks(key, i.Value, cb)
+			is.runCallbacks(key, i.Value, i.OrigStamp, cb)
 		}
 		return nil
 	}, true /* deleteExpired */); err != nil {
@@ -446,20 +449,24 @@ func (is *infoStore) registerCallback(
 // processCallbacks processes callbacks for the specified key by
 // matching each callback's regular expression against the key and invoking
 // the corresponding callback method on a match.
-func (is *infoStore) processCallbacks(key string, content roachpb.Value, changed bool) {
+func (is *infoStore) processCallbacks(
+	key string, content roachpb.Value, origTimestamp int64, changed bool,
+) {
 	var callbacks []*callback
 	for _, cb := range is.callbacks {
 		if (changed || cb.redundant) && cb.matcher.MatchString(key) {
 			callbacks = append(callbacks, cb)
 		}
 	}
-	is.runCallbacks(key, content, callbacks...)
+	is.runCallbacks(key, content, origTimestamp, callbacks...)
 }
 
 // runCallbacks receives a list of callbacks and contents that match the key.
 // It adds work to the callback work slices, and signals the associated callback
 // workers to execute the work.
-func (is *infoStore) runCallbacks(key string, content roachpb.Value, callbacks ...*callback) {
+func (is *infoStore) runCallbacks(
+	key string, content roachpb.Value, origTimestamp int64, callbacks ...*callback,
+) {
 	// Check if the stopper is quiescing. If so, do not add the callbacks to the
 	// callback work list because they won't be processed anyways.
 	select {
@@ -478,6 +485,7 @@ func (is *infoStore) runCallbacks(key string, content roachpb.Value, callbacks .
 			method:         cb.method,
 			key:            key,
 			content:        content,
+			origTimestamp:  origTimestamp,
 		})
 		cb.cw.mu.Unlock()
 

pkg/gossip/infostore_test.go

@@ -479,32 +479,48 @@ func TestLeastUseful(t *testing.T) {
 }
 
 type callbackRecord struct {
-	keys []string
-	wg   *sync.WaitGroup
+	key       string
+	value     roachpb.Value
+	timestamp int64
+}
+
+type callbackRecords struct {
+	records []callbackRecord
+	wg      *sync.WaitGroup
 	syncutil.Mutex
 }
 
-func (cr *callbackRecord) Add(key string, _ roachpb.Value) {
+func (cr *callbackRecords) Add(key string, value roachpb.Value, origTs int64) {
 	cr.Lock()
 	defer cr.Unlock()
-	cr.keys = append(cr.keys, key)
+	cr.records = append(cr.records, callbackRecord{key, value, origTs})
 	cr.wg.Done()
 }
 
-func (cr *callbackRecord) Keys() []string {
+func (cr *callbackRecords) Keys() []string {
 	cr.Lock()
 	defer cr.Unlock()
-	return append([]string(nil), cr.keys...)
+	keys := make([]string, len(cr.records))
+	for i, record := range cr.records {
+		keys[i] = record.key
+	}
+	return keys
+}
+
+func (cr *callbackRecords) Records() []callbackRecord {
+	cr.Lock()
+	defer cr.Unlock()
+	return append([]callbackRecord(nil), cr.records...)
 }
 
 func TestCallbacks(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	is, stopper := newTestInfoStore()
 	defer stopper.Stop(context.Background())
 	wg := &sync.WaitGroup{}
-	cb1 := callbackRecord{wg: wg}
-	cb2 := callbackRecord{wg: wg}
-	cbAll := callbackRecord{wg: wg}
+	cb1 := callbackRecords{wg: wg}
+	cb2 := callbackRecords{wg: wg}
+	cbAll := callbackRecords{wg: wg}
 
 	unregisterCB1 := is.registerCallback("key1", cb1.Add)
 	is.registerCallback("key2", cb2.Add)
@@ -603,18 +619,17 @@ func TestCallbacks(t *testing.T) {
 	}
 }
 
-// TestRegisterCallback verifies that a callback is invoked when
-// registered if there are items which match its regexp in the
-// infostore.
+// TestRegisterCallback verifies that a callback is invoked correctly when
+// registered if there are items which match its regexp in the infostore.
 func TestRegisterCallback(t *testing.T) {
 	defer leaktest.AfterTest(t)()
 	is, stopper := newTestInfoStore()
 	defer stopper.Stop(context.Background())
 	wg := &sync.WaitGroup{}
-	cb := callbackRecord{wg: wg}
+	cb := callbackRecords{wg: wg}
 
-	i1 := is.newInfo(nil, time.Second)
-	i2 := is.newInfo(nil, time.Second)
+	i1 := is.newInfo([]byte("val1"), time.Second)
+	i2 := is.newInfo([]byte("val2"), time.Second)
 	if err := is.addInfo("key1", i1); err != nil {
 		t.Fatal(err)
 	}
@@ -623,13 +638,35 @@ func TestRegisterCallback(t *testing.T) {
 	}
 
 	wg.Add(2)
+	// Register a callback after the infos are added.
 	is.registerCallback("key.*", cb.Add)
 	wg.Wait()
-	actKeys := cb.Keys()
-	sort.Strings(actKeys)
-	if expKeys := []string{"key1", "key2"}; !reflect.DeepEqual(actKeys, expKeys) {
-		t.Errorf("expected %v, got %v", expKeys, cb.Keys())
+	actRecords := cb.Records()
+	// Sort records by key since callback order is not guaranteed.
+	sort.Slice(actRecords, func(i, j int) bool {
+		return actRecords[i].key < actRecords[j].key
+	})
+
+	expectedRecords := []callbackRecord{
+		{key: "key1", value: i1.Value, timestamp: i1.OrigStamp},
+		{key: "key2", value: i2.Value, timestamp: i2.OrigStamp},
+	}
+
+	require.Equal(t, expectedRecords, actRecords)
+
+	// Verify callback fires for new matching info
+	i3 := is.newInfo([]byte("val3"), time.Second)
+	wg.Add(1)
+	if err := is.addInfo("key3", i3); err != nil {
+		t.Fatal(err)
 	}
+	wg.Wait()
+	actRecords = cb.Records()
+	sort.Slice(actRecords, func(i, j int) bool {
+		return actRecords[i].key < actRecords[j].key
+	})
+	expectedRecords = append(expectedRecords, callbackRecord{key: "key3", value: i3.Value, timestamp: i3.OrigStamp})
+	require.Equal(t, expectedRecords, actRecords)
 }
 
 // TestCallbacksCalledSequentially verifies that callbacks are called in a
@@ -650,14 +687,14 @@ func TestCallbacksCalledSequentially(t *testing.T) {
 
 	// Create a callback generator that will generate a callback that will
 	// assert that the keys are in sequential order.
-	callbackGenerator := func() func(key string, _ roachpb.Value) {
+	callbackGenerator := func() func(key string, _ roachpb.Value, _ int64) {
 		// Add the number of updates to the wait group.
 		wg.Add(numUpdates)
 		// Initially, the expected next key is 0.
 		expectedNextKey := 0
 
 		// Return a callback that will assert the key is in sequential order.
-		return func(key string, _ roachpb.Value) {
+		return func(key string, _ roachpb.Value, _ int64) {
 			// Convert key to int and assert it matches the expected value.
 			keyInt, err := strconv.Atoi(key)
 			require.NoError(t, err)
@@ -672,7 +709,7 @@ func TestCallbacksCalledSequentially(t *testing.T) {
 	// halfway through the test to assert that it doesn't impact the
 	// sequential order of the other callbacks.
 	is.registerCallback(".*", callbackGenerator())
-	unregister := is.registerCallback(".*", func(key string, _ roachpb.Value) {})
+	unregister := is.registerCallback(".*", func(key string, _ roachpb.Value, _ int64) {})
 	is.registerCallback(".*", callbackGenerator())
 
 	for i := range numUpdates {
@@ -694,7 +731,7 @@ func BenchmarkCallbackParallelism(b *testing.B) {
 	defer stopper.Stop(ctx)
 	wg := &sync.WaitGroup{}
 
-	callback := func(key string, val roachpb.Value) {
+	callback := func(key string, val roachpb.Value, _ int64) {
 		// Sleep for a short duration to simulate work done in callback.
 		time.Sleep(time.Millisecond)
 		wg.Done()

pkg/gossip/server.go

@@ -396,7 +396,7 @@ func (s *server) start(addr net.Addr) {
 	// We require redundant callbacks here as the broadcast callback is
 	// propagating gossip infos to other nodes and needs to propagate the new
 	// expiration info.
-	unregister := s.mu.is.registerCallback(".*", func(_ string, _ roachpb.Value) {
+	unregister := s.mu.is.registerCallback(".*", func(_ string, _ roachpb.Value, _ int64) {
 		broadcast()
 	}, Redundant)
 

pkg/kv/kvserver/allocator/allocatorimpl/allocator_test.go

@@ -8829,7 +8829,7 @@ func TestAllocatorFullDisks(t *testing.T) {
 
 	var wg sync.WaitGroup
 	g.RegisterCallback(gossip.MakePrefixPattern(gossip.KeyStoreDescPrefix),
-		func(_ string, _ roachpb.Value) { wg.Done() },
+		func(_ string, _ roachpb.Value, _ int64) { wg.Done() },
 		// Redundant callbacks are required by this test.
 		gossip.Redundant)
 

pkg/kv/kvserver/allocator/storepool/store_pool.go

@@ -482,7 +482,7 @@ func (sp *StorePool) statusString(nl NodeLivenessFunc) redact.RedactableString {
 }
 
 // storeGossipUpdate is the Gossip callback used to keep the StorePool up to date.
-func (sp *StorePool) storeGossipUpdate(_ string, content roachpb.Value) {
+func (sp *StorePool) storeGossipUpdate(_ string, content roachpb.Value, _ int64) {
 	var storeDesc roachpb.StoreDescriptor
 
 	if err := content.GetProto(&storeDesc); err != nil {

pkg/kv/kvserver/client_split_test.go

@@ -2817,7 +2817,7 @@ func TestStoreRangeGossipOnSplits(t *testing.T) {
 
 	var lastSD roachpb.StoreDescriptor
 	rangeCountCh := make(chan int32)
-	unregister := store.Gossip().RegisterCallback(storeKey, func(_ string, val roachpb.Value) {
+	unregister := store.Gossip().RegisterCallback(storeKey, func(_ string, val roachpb.Value, _ int64) {
 		var sd roachpb.StoreDescriptor
 		if err := val.GetProto(&sd); err != nil {
 			panic(err)

pkg/kv/kvserver/gossip_test.go

@@ -40,7 +40,7 @@ func TestGossipFirstRange(t *testing.T) {
 	descs := make(chan *roachpb.RangeDescriptor)
 	unregister := tc.Servers[0].GossipI().(*gossip.Gossip).
 		RegisterCallback(gossip.KeyFirstRangeDescriptor,
-			func(_ string, content roachpb.Value) {
+			func(_ string, content roachpb.Value, _ int64) {
 				var desc roachpb.RangeDescriptor
 				if err := content.GetProto(&desc); err != nil {
 					select {