fix(bigtable/bttest): make table gc release memory

bigtable/bttest/inmem.go

@@ -1440,6 +1440,27 @@ func (t *table) mutableRow(key string) *row {
 }
 
 func (t *table) gc() {
+	toDelete := t.gcReadOnly()
+	if len(toDelete) == 0 {
+		return
+	}
+
+	// We delete rows that no longer have any cells
+	t.mu.Lock()
+	defer t.mu.Unlock()
+	for _, i := range toDelete {
+		r := i.(*row)
+		// Make sure the row still has no cells. We've not been holding a lock
+		// so it could have changed since we checked it.
+		r.mu.Lock()
+		if len(r.families) == 0 {
+			t.rows.Delete(i)
+		}
+		r.mu.Unlock()
+	}
+}
+
+func (t *table) gcReadOnly() (toDelete []btree.Item) {
 	// This method doesn't add or remove rows, so we only need a read lock for the table.
 	t.mu.RLock()
 	defer t.mu.RUnlock()
@@ -1452,16 +1473,23 @@ func (t *table) gc() {
 		}
 	}
 	if len(rules) == 0 {
-		return
+		return nil
 	}
 
+	// It isn't clear whether it's safe to delete within the iterator, so we do
+	// not
 	t.rows.Ascend(func(i btree.Item) bool {
 		r := i.(*row)
 		r.mu.Lock()
 		r.gc(rules)
+		if len(r.families) == 0 {
+			toDelete = append(toDelete, i)
+		}
 		r.mu.Unlock()
 		return true
 	})
+
+	return toDelete
 }
 
 type byRowKey []*row
@@ -1547,7 +1575,15 @@ func (r *row) gc(rules map[string]*btapb.GcRule) {
 			continue
 		}
 		for col, cs := range fam.cells {
-			r.families[fam.name].cells[col] = applyGC(cs, rule)
+			cs = applyGC(cs, rule)
+			if len(cs) == 0 {
+				delete(fam.cells, col)
+			} else {
+				fam.cells[col] = cs
+			}
+		}
+		if len(fam.cells) == 0 {
+			delete(r.families, fam.name)
 		}
 	}
 }

bigtable/bttest/inmem_test.go

@@ -40,6 +40,7 @@ import (
 	"google.golang.org/protobuf/encoding/prototext"
 	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/testing/protocmp"
+	"google.golang.org/protobuf/types/known/durationpb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
 )
 
@@ -472,6 +473,72 @@ func TestModifyColumnFamilies(t *testing.T) {
 	readRows(18, 6, 2)
 }
 
+func TestGC(t *testing.T) {
+	// Create server
+	s := &server{
+		tables: make(map[string]*table),
+	}
+	ctx := context.Background()
+
+	colFamilyID := "colFam"
+	colName := "colName"
+	rowKey := "rowKey"
+
+	// Create table with max age gc rule
+	newTbl := btapb.Table{
+		ColumnFamilies: map[string]*btapb.ColumnFamily{},
+	}
+	newTbl.ColumnFamilies[colFamilyID] = &btapb.ColumnFamily{GcRule: &btapb.GcRule{Rule: &btapb.GcRule_MaxAge{MaxAge: durationpb.New(time.Millisecond)}}}
+
+	tblInfo, err := s.CreateTable(ctx, &btapb.CreateTableRequest{Parent: "cluster", TableId: "t", Table: &newTbl})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Populate the table
+	for i := 0; i < 2; i++ {
+		req := &btpb.MutateRowRequest{
+			TableName: tblInfo.Name,
+			RowKey:    []byte(rowKey),
+			Mutations: []*btpb.Mutation{{
+				Mutation: &btpb.Mutation_SetCell_{
+					SetCell: &btpb.Mutation_SetCell{
+						FamilyName:      colFamilyID,
+						ColumnQualifier: []byte(colName),
+						TimestampMicros: 1000, // MaxAge is 1ms
+						Value:           []byte{},
+					}},
+			}},
+		}
+		if _, err := s.MutateRow(ctx, req); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// Sleep till maxAge passes
+	time.Sleep(2 * time.Millisecond)
+
+	// Trigger gc
+	tbl := s.tables[tblInfo.Name]
+	tbl.gc()
+
+	// Verify that the row was deleted after garbage collection
+	readRowsReq := &btpb.ReadRowsRequest{
+		TableName: tblInfo.Name,
+		Rows:      &btpb.RowSet{RowKeys: [][]byte{[]byte(rowKey)}},
+	}
+	mock := &MockReadRowsServer{}
+	if err = s.ReadRows(readRowsReq, mock); err != nil {
+		t.Errorf("ReadRows error: %v", err)
+	}
+	if got, want := len(mock.responses), 0; got != want {
+		t.Errorf("response count: got %d, want %d", got, want)
+	}
+}
+
 func TestDropRowRange(t *testing.T) {
 	s := &server{
 		tables: make(map[string]*table),