tan: fixed getEntries when size is limited

binbinshi · Aug 3, 2022 · a83f853 · a83f853
1 parent e77cf6b
commit a83f853
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 0 deletions.
diff --git a/internal/tan/db.go b/internal/tan/db.go
@@ -276,6 +276,7 @@ func (d *db) getEntries(shardID uint64, replicaID uint64,
 		maxSize = math.MaxUint64
 	}
 	expected := low
+	done := false
 	for _, ie := range ies {
 		queryIndex := ie
 		f := func(u pb.Update, _ int64) bool {
@@ -294,6 +295,7 @@ func (d *db) getEntries(shardID uint64, replicaID uint64,
 					}
 					entries = append(entries, e)
 					if size > maxSize {
+						done = true
 						return false
 					}
 				} else {
@@ -305,6 +307,9 @@ func (d *db) getEntries(shardID uint64, replicaID uint64,
 		if err := d.readLog(ie, f); err != nil {
 			return nil, 0, err
 		}
+		if done {
+			return entries, size, nil
+		}
 	}
 	return entries, size, nil
 }

diff --git a/internal/tan/db_test.go b/internal/tan/db_test.go
@@ -665,3 +665,31 @@ func TestRebuildLog(t *testing.T) {
 	require.NoError(t, err)
 	require.Equal(t, fi.Size()*19/20, fi2.Size())
 }
+
+func TestGetEntriesWithMaxSize(t *testing.T) {
+	fs := vfs.NewMem()
+	opts := &Options{
+		MaxManifestFileSize: MaxManifestFileSize,
+		MaxLogFileSize:      1024,
+		FS:                  fs,
+	}
+	tf := func(t *testing.T, db *db) {
+		cmd := make([]byte, 128)
+		buf := make([]byte, 1024)
+		for i := 0; i < 128; i++ {
+			u := pb.Update{
+				ShardID:   2,
+				ReplicaID: 3,
+				EntriesToSave: []pb.Entry{
+					{Index: 1 + uint64(i), Term: 5, Cmd: cmd},
+				},
+			}
+			_, err := db.write(u, buf)
+			require.NoError(t, err)
+		}
+		entries, _, err := db.getEntries(2, 3, nil, 0, 1, 128, 128)
+		require.NoError(t, err)
+		require.Equal(t, 1, len(entries))
+	}
+	runTanTest(t, opts, tf, fs)
+}