Include head & commit offset in GetStatus (oxia-db#326)

Author: merlimat

proto/replication.pb.go

@@ -145,6 +145,6 @@ message GetStatusResponse {
   int64 term = 1;
   ServingStatus status = 2;
 
-  // More info to add for leaders
-  // ensemble, lag for each follower, etc...
+  int64 head_offset = 3;
+  int64 commit_offset = 4;
 }

proto/replication.proto

@@ -662,8 +662,10 @@ func (fc *followerController) GetStatus(request *proto.GetStatusRequest) (*proto
 	defer fc.Unlock()
 
 	return &proto.GetStatusResponse{
-		Term:   fc.term,
-		Status: fc.status,
+		Term:         fc.term,
+		Status:       fc.status,
+		HeadOffset:   fc.wal.LastOffset(),
+		CommitOffset: fc.CommitOffset(),
 	}, nil
 }
 

server/follower_controller.go

@@ -820,6 +820,53 @@ func TestFollowerController_Closed(t *testing.T) {
 	assert.NoError(t, walFactory.Close())
 }
 
+func TestFollower_GetStatus(t *testing.T) {
+	var shardId int64
+	kvFactory, _ := kv.NewPebbleKVFactory(testKVOptions)
+	walFactory := wal.NewInMemoryWalFactory()
+
+	fc, _ := NewFollowerController(Config{}, common.DefaultNamespace, shardId, walFactory, kvFactory)
+	_, _ = fc.NewTerm(&proto.NewTermRequest{Term: 2})
+
+	stream := newMockServerReplicateStream()
+	go func() {
+		//cancelled due to fc.Close() below
+		assert.ErrorIs(t, fc.Replicate(stream), context.Canceled)
+	}()
+
+	stream.AddRequest(createAddRequest(t, 2, 0, map[string]string{"a": "0", "b": "1"}, wal.InvalidOffset))
+	stream.AddRequest(createAddRequest(t, 2, 1, map[string]string{"a": "0", "b": "1"}, 0))
+	stream.AddRequest(createAddRequest(t, 2, 2, map[string]string{"a": "0", "b": "1"}, 1))
+
+	// Wait for responses
+	r1 := stream.GetResponse()
+	assert.EqualValues(t, 0, r1.Offset)
+
+	r2 := stream.GetResponse()
+	assert.EqualValues(t, 1, r2.Offset)
+
+	r3 := stream.GetResponse()
+	assert.EqualValues(t, 2, r3.Offset)
+
+	assert.Eventually(t, func() bool {
+		res, _ := fc.GetStatus(&proto.GetStatusRequest{ShardId: shardId})
+		return res.CommitOffset == 1
+	}, 10*time.Second, 100*time.Millisecond)
+
+	res, err := fc.GetStatus(&proto.GetStatusRequest{ShardId: shardId})
+	assert.NoError(t, err)
+	assert.Equal(t, &proto.GetStatusResponse{
+		Term:         2,
+		Status:       proto.ServingStatus_FOLLOWER,
+		HeadOffset:   2,
+		CommitOffset: 1,
+	}, res)
+
+	assert.NoError(t, fc.Close())
+	assert.NoError(t, kvFactory.Close())
+	assert.NoError(t, walFactory.Close())
+}
+
 func closeChanIsNotNil(fc FollowerController) func() bool {
 	return func() bool {
 		_fc := fc.(*followerController)

server/follower_controller_test.go

@@ -859,8 +859,10 @@ func (lc *leaderController) GetStatus(request *proto.GetStatusRequest) (*proto.G
 	defer lc.RUnlock()
 
 	return &proto.GetStatusResponse{
-		Term:   lc.term,
-		Status: lc.status,
+		Term:         lc.term,
+		Status:       lc.status,
+		HeadOffset:   lc.quorumAckTracker.HeadOffset(),
+		CommitOffset: lc.quorumAckTracker.CommitOffset(),
 	}, nil
 }
 

server/leader_controller.go

@@ -996,3 +996,46 @@ func TestLeaderController_DeleteShard(t *testing.T) {
 	assert.NoError(t, lc.Close())
 	assert.NoError(t, walFactory.Close())
 }
+
+func TestLeaderController_GetStatus(t *testing.T) {
+	var shard int64 = 1
+
+	kvFactory, _ := kv.NewPebbleKVFactory(testKVOptions)
+	walFactory := wal.NewInMemoryWalFactory()
+
+	lc, _ := NewLeaderController(Config{}, common.DefaultNamespace, shard, newMockRpcClient(), walFactory, kvFactory)
+	_, _ = lc.NewTerm(&proto.NewTermRequest{ShardId: shard, Term: 2})
+	_, _ = lc.BecomeLeader(&proto.BecomeLeaderRequest{
+		ShardId:           shard,
+		Term:              2,
+		ReplicationFactor: 1,
+		FollowerMaps:      nil,
+	})
+
+	/// Write entry
+	_, _ = lc.Write(&proto.WriteRequest{
+		ShardId: &shard,
+		Puts: []*proto.PutRequest{{
+			Key:   "a",
+			Value: []byte("value-a")}},
+	})
+	_, _ = lc.Write(&proto.WriteRequest{
+		ShardId: &shard,
+		Puts: []*proto.PutRequest{{
+			Key:   "b",
+			Value: []byte("value-b")}},
+	})
+
+	res, err := lc.GetStatus(&proto.GetStatusRequest{ShardId: shard})
+	assert.NoError(t, err)
+	assert.Equal(t, &proto.GetStatusResponse{
+		Term:         2,
+		Status:       proto.ServingStatus_LEADER,
+		HeadOffset:   1,
+		CommitOffset: 1,
+	}, res)
+
+	assert.NoError(t, lc.Close())
+	assert.NoError(t, kvFactory.Close())
+	assert.NoError(t, walFactory.Close())
+}