Repository: polarsignals/frostdb (go)

Issue:
Recycle goroutine doing compaction
While working on https://github.com/polarsignals/frostdb/pull/118 I realised that we create a go-routine per compaction for each granule.

I think this should be bounded but also we should probably recycle go-routine since those can grow pretty large in term of memory

<img width="2487" alt="image" src="https://user-images.githubusercontent.com/1053421/177159199-6c105754-d2a7-4bef-9d0c-f69438d8cad5.png">


I suggest we create 8 go-routines that will take care of compactions, then send granule to compact via  a buffered channel of 32, meaning it can block on the write path.

Pull request description:
frostdb: add compactorPool to limit compaction goroutines
This commit adds a worker pool of goroutines (`runtime.NumCPU` by default) per database that will iterate over all the tables in that database and attempt to perform a compaction every defined sweep interval (`100ms` by default).

The compaction logic has not been modified.

This PR also adds a `compacting` atomic boolean to the granule so that no two goroutines attempt to compact the same granule. Previously, I think this was the case because the only atomic that was checked was `pruned`, which is only set once the granule has been pruned from the index.

Closes #119 

Patch:
diff --git a/compaction.go b/compaction.go
new file mode 100644
index 0000000..12fe94f
--- /dev/null
+++ b/compaction.go
@@ -0,0 +1,109 @@
+package frostdb
+
+import (
+	"context"
+	"runtime"
+	"sync"
+	"time"
+
+	"github.com/go-kit/log/level"
+	"github.com/google/btree"
+	"github.com/thanos-io/objstore/errutil"
+)
+
+type compactorPool struct {
+	db *DB
+
+	concurrency int
+	interval    time.Duration
+	wg          sync.WaitGroup
+	cancel      context.CancelFunc
+}
+
+var (
+	defaultConcurrency   = runtime.GOMAXPROCS(0)
+	defaultSweepInterval = 100 * time.Millisecond
+)
+
+func newCompactorPool(db *DB) *compactorPool {
+	concurrency := defaultConcurrency
+	if override := db.columnStore.compactionOptions.concurrency; override != 0 {
+		concurrency = override
+	}
+
+	sweepInterval := defaultSweepInterval
+	if override := db.columnStore.compactionOptions.sweepInterval; override != 0 {
+		sweepInterval = override
+	}
+	return &compactorPool{
+		db:          db,
+		concurrency: concurrency,
+		interval:    sweepInterval,
+	}
+}
+
+func (c *compactorPool) start() {
+	ctx, cancelFn := context.WithCancel(context.Background())
+	c.cancel = cancelFn
+	for i := 0; i < c.concurrency; i++ {
+		c.wg.Add(1)
+		go func() {
+			defer c.wg.Done()
+			c.compactLoop(ctx)
+		}()
+	}
+}
+
+func (c *compactorPool) stop() {
+	if c.cancel == nil {
+		// Pool was not started.
+		return
+	}
+	c.cancel()
+	c.wg.Wait()
+}
+
+func (c *compactorPool) compactLoop(ctx context.Context) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-time.After(c.interval):
+			c.db.mtx.RLock()
+			tablesToCompact := make([]*Table, 0, len(c.db.tables))
+			for _, table := range c.db.tables {
+				tablesToCompact = append(tablesToCompact, table)
+			}
+			c.db.mtx.RUnlock()
+
+			for _, table := range tablesToCompact {
+				if err := table.ActiveBlock().compact(); err != nil {
+					level.Warn(c.db.logger).Log("msg", "compaction failed", "err", err)
+				}
+			}
+		}
+	}
+}
+
+func (t *TableBlock) compact() error {
+	var compactionErrors errutil.MultiError
+	t.Index().Ascend(func(i btree.Item) bool {
+		granuleToCompact := i.(*Granule)
+		if granuleToCompact.metadata.size.Load() < uint64(t.table.db.columnStore.granuleSizeBytes) {
+			// Skip granule since its size is under the target size.
+			return true
+		}
+		if !granuleToCompact.metadata.pruned.CompareAndSwap(0, 1) {
+			// Someone else claimed this granule compaction.
+			return true
+		}
+
+		defer t.table.metrics.compactions.Inc()
+		if err := t.compactGranule(granuleToCompact); err != nil {
+			t.abort(granuleToCompact)
+			compactionErrors.Add(err)
+		}
+		return true
+	})
+	return compactionErrors.Err()
+}
diff --git a/db.go b/db.go
index f055fc1..6ee3b11 100644
--- a/db.go
+++ b/db.go
@@ -40,6 +40,10 @@ type ColumnStore struct {
 	bucket               storage.Bucket
 	ignoreStorageOnQuery bool
 	enableWAL            bool
+	compactionOptions    struct {
+		concurrency   int
+		sweepInterval time.Duration
+	}
 
 	// indexDegree is the degree of the btree index (default = 2)
 	indexDegree int
@@ -162,6 +166,22 @@ func WithIgnoreStorageOnQuery() Option {
 	}
 }
 
+// WithCompactionConcurrency specfies the number of concurrent goroutines
+// compacting data for each database.
+func WithCompactionConcurrency(c int) Option {
+	return func(s *ColumnStore) error {
+		s.compactionOptions.concurrency = c
+		return nil
+	}
+}
+
+func WithCompactionSweepInterval(interval time.Duration) Option {
+	return func(s *ColumnStore) error {
+		s.compactionOptions.sweepInterval = interval
+		return nil
+	}
+}
+
 // Close persists all data from the columnstore to storage.
 // It is no longer valid to use the coumnstore for reads or writes, and the object should not longer be reused.
 func (s *ColumnStore) Close() error {
@@ -241,6 +261,8 @@ type DB struct {
 	// TxPool is a waiting area for finished transactions that haven't been added to the watermark
 	txPool *TxPool
 
+	compactorPool *compactorPool
+
 	// highWatermark maintains the highest consecutively completed tx number
 	highWatermark *atomic.Uint64
 
@@ -300,6 +322,7 @@ func (s *ColumnStore) DB(ctx context.Context, name string) (*DB, error) {
 
 	if dbSetupErr := func() error {
 		db.txPool = NewTxPool(db.highWatermark)
+		db.compactorPool = newCompactorPool(db)
 		// If bucket storage is configured; scan for existing tables in the database
 		if db.bucket != nil {
 			if err := db.bucket.Iter(ctx, "", func(tableName string) error {
@@ -330,6 +353,7 @@ func (s *ColumnStore) DB(ctx context.Context, name string) (*DB, error) {
 		return nil, dbSetupErr
 	}
 
+	db.compactorPool.start()
 	s.dbs[name] = db
 	return db, nil
 }
@@ -412,12 +436,11 @@ func (db *DB) replayWAL(ctx context.Context) error {
 					return fmt.Errorf("instantiate table: %w", err)
 				}
 
-				db.tables[tableName] = table
-
 				table.active, err = newTableBlock(table, 0, tx, id)
 				if err != nil {
 					return err
 				}
+				db.tables[tableName] = table
 				return nil
 			}
 			if err != nil {
@@ -507,6 +530,10 @@ func (db *DB) Close() error {
 		db.txPool.Stop()
 	}
 
+	if db.compactorPool != nil {
+		db.compactorPool.stop()
+	}
+
 	return nil
 }
 
diff --git a/granule.go b/granule.go
index a44b3f8..6cba12a 100644
--- a/granule.go
+++ b/granule.go
@@ -28,13 +28,15 @@ type Granule struct {
 // GranuleMetadata is the metadata for a granule.
 type GranuleMetadata struct {
 	// least is the row that exists within the Granule that is the least.
-	// This is used for quick insertion into the btree, without requiring an iterator
+	// This is used for quick insertion into the btree, without requiring an
+	// iterator.
 	least atomic.Pointer[dynparquet.DynamicRow]
 
-	// size is the raw commited, and uncommited size of the granule. It is used as a suggestion for potential compaction
+	// size is the raw committed, and uncommitted size of the granule. It is
+	// used as a suggestion for potential compaction.
 	size *atomic.Uint64
 
-	// pruned indicates if this Granule is longer found in the index
+	// pruned indicates if this Granule is longer found in the index.
 	pruned *atomic.Uint64
 }
 
diff --git a/table.go b/table.go
index 00da2f3..e3e9f5c 100644
--- a/table.go
+++ b/table.go
@@ -139,7 +139,6 @@ type TableBlock struct {
 
 	pendingWritersWg sync.WaitGroup
 
-	wg  *sync.WaitGroup
 	mtx *sync.RWMutex
 }
 
@@ -280,7 +279,6 @@ func (t *Table) newTableBlock(prevTx, tx uint64, id ulid.ULID) error {
 func (t *Table) writeBlock(block *TableBlock) {
 	level.Debug(t.logger).Log("msg", "syncing block")
 	block.pendingWritersWg.Wait()
-	block.wg.Wait()
 
 	// from now on, the block will no longer be modified, we can persist it to disk
 
@@ -392,8 +390,8 @@ func (t *Table) Schema() *dynparquet.Schema {
 	return t.config.schema
 }
 
-func (t *Table) Sync() {
-	t.ActiveBlock().Sync()
+func (t *Table) EnsureCompaction() error {
+	return t.ActiveBlock().EnsureCompaction()
 }
 
 // Write objects into the table.
@@ -912,7 +910,6 @@ func newTableBlock(table *Table, prevTx, tx uint64, id ulid.ULID) (*TableBlock,
 	tb := &TableBlock{
 		table:  table,
 		index:  &index,
-		wg:     &sync.WaitGroup{},
 		mtx:    &sync.RWMutex{},
 		ulid:   id,
 		size:   &atomic.Int64{},
@@ -930,10 +927,9 @@ func newTableBlock(table *Table, prevTx, tx uint64, id ulid.ULID) (*TableBlock,
 	return tb, nil
 }
 
-// Sync the table block. This will return once all writes have completed and
-// all potentially started split operations have completed.
-func (t *TableBlock) Sync() {
-	t.wg.Wait()
+// EnsureCompaction forces a TableBlock compaction.
+func (t *TableBlock) EnsureCompaction() error {
+	return t.compact()
 }
 
 func (t *TableBlock) Insert(ctx context.Context, tx uint64, buf *dynparquet.SerializedBuffer) error {
@@ -999,15 +995,10 @@ func (t *TableBlock) Insert(ctx context.Context, tx uint64, buf *dynparquet.Seri
 			}
 
 			part := NewPart(tx, serBuf)
-			size, err := granule.AddPart(part)
-			if err != nil {
+			if _, err := granule.AddPart(part); err != nil {
 				return fmt.Errorf("failed to add part to granule: %w", err)
 			}
 			parts = append(parts, part)
-			if size >= uint64(t.table.db.columnStore.granuleSizeBytes) {
-				t.wg.Add(1)
-				go t.compact(granule)
-			}
 			t.size.Add(serBuf.ParquetFile().Size())
 
 			b = bytes.NewBuffer(nil)
@@ -1300,16 +1291,6 @@ func (t *TableBlock) splitRowsByGranule(buf *dynparquet.SerializedBuffer) (map[*
 	return rowsByGranule, nil
 }
 
-// compact will compact a Granule; should be performed as a background go routine.
-func (t *TableBlock) compact(g *Granule) {
-	defer t.wg.Done()
-	if err := t.compactGranule(g); err != nil {
-		t.abort(g)
-		level.Error(t.logger).Log("msg", "failed to compact granule", "err", err)
-	}
-	t.table.metrics.compactions.Inc()
-}
-
 // addPartToGranule finds the corresponding granule it belongs to in a sorted list of Granules.
 func addPartToGranule(granules []*Granule, p *Part) error {
 	row, err := p.Least()


Test patch:
diff --git a/aggregate_test.go b/aggregate_test.go
index 36488c7..0d0c62f 100644
--- a/aggregate_test.go
+++ b/aggregate_test.go
@@ -82,9 +82,6 @@ func TestAggregate(t *testing.T) {
 	_, err = table.InsertBuffer(context.Background(), buf)
 	require.NoError(t, err)
 
-	// Ensure all transactions are completed
-	table.Sync()
-
 	engine := query.NewEngine(
 		memory.NewGoAllocator(),
 		db.TableProvider(),
diff --git a/bench_test.go b/bench_test.go
index 2918528..d0a725f 100644
--- a/bench_test.go
+++ b/bench_test.go
@@ -53,7 +53,9 @@ func newDBForBenchmarks(ctx context.Context, b testing.TB) (*ColumnStore, *DB, e
 	if err != nil {
 		return nil, nil, err
 	}
-	table.Sync()
+	if err := table.EnsureCompaction(); err != nil {
+		return nil, nil, err
+	}
 
 	b.Logf("db initialized and WAL replayed, starting benchmark %s", b.Name())
 	return col, colDB, nil
diff --git a/compaction_test.go b/compaction_test.go
new file mode 100644
index 0000000..e3b2c91
--- /dev/null
+++ b/compaction_test.go
@@ -0,0 +1,83 @@
+package frostdb
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/uuid"
+	"github.com/stretchr/testify/require"
+
+	"github.com/polarsignals/frostdb/dynparquet"
+)
+
+func TestCompaction(t *testing.T) {
+	c, table := basicTable(t)
+	defer c.Close()
+	table.config.rowGroupSize = 2 // override row group size setting
+
+	samples := dynparquet.Samples{{
+		Labels: []dynparquet.Label{
+			{Name: "label1", Value: "value1"},
+			{Name: "label2", Value: "value2"},
+		},
+		Stacktrace: []uuid.UUID{
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
+		},
+		Timestamp: 1,
+		Value:     1,
+	}, {
+		Labels: []dynparquet.Label{
+			{Name: "label1", Value: "value1"},
+			{Name: "label2", Value: "value2"},
+			{Name: "label3", Value: "value3"},
+		},
+		Stacktrace: []uuid.UUID{
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
+		},
+		Timestamp: 2,
+		Value:     2,
+	}, {
+		Labels: []dynparquet.Label{
+			{Name: "label1", Value: "value1"},
+			{Name: "label2", Value: "value2"},
+			{Name: "label4", Value: "value4"},
+		},
+		Stacktrace: []uuid.UUID{
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
+			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
+		},
+		Timestamp: 3,
+		Value:     3,
+	}}
+
+	ctx := context.Background()
+	tx, err := table.Write(ctx, samples[0], samples[1], samples[2])
+	require.NoError(t, err)
+
+	// Wait for the writes to have been marked as completed
+	table.db.Wait(tx)
+
+	// Expect only a single granule to have been created
+	require.Equal(t, 1, table.active.Index().Len())
+
+	// Force compaction on the granule
+	g := (table.active.Index().Min()).(*Granule)
+	table.active.compactGranule(g)
+
+	// Expect the granule to have been compacted into a single granule
+	require.Equal(t, 1, table.active.Index().Len())
+
+	// Expect the granule to contain a single part with multiple row groups
+	g = (table.active.Index().Min()).(*Granule)
+	parts := 0
+	g.parts.Iterate(func(p *Part) bool {
+		parts++
+		require.Equal(t, 2, p.Buf.NumRowGroups())
+		return true
+	})
+
+	// There should only be a single part that contains multiple row groups
+	require.Equal(t, 1, parts)
+}
diff --git a/db_test.go b/db_test.go
index 9aa8f4d..2ef7881 100644
--- a/db_test.go
+++ b/db_test.go
@@ -64,7 +64,7 @@ func TestDBWithWALAndBucket(t *testing.T) {
 		_, err = table.InsertBuffer(ctx, buf)
 		require.NoError(t, err)
 	}
-	table.Sync()
+	require.NoError(t, table.EnsureCompaction())
 	require.NoError(t, c.Close())
 
 	c, err = New(
diff --git a/table_test.go b/table_test.go
index e454987..1c7b4b8 100644
--- a/table_test.go
+++ b/table_test.go
@@ -413,8 +413,7 @@ func benchmarkTableInserts(b *testing.B, rows, iterations, writers int) {
 
 		b.StopTimer()
 		pool := memory.NewGoAllocator()
-		// Wait for all compaction routines to complete
-		table.Sync()
+		require.NoError(b, table.EnsureCompaction())
 
 		// Calculate the number of entries in database
 		totalrows := int64(0)
@@ -1321,76 +1320,3 @@ func Test_DoubleTable(t *testing.T) {
 	})
 	require.NoError(t, err)
 }
-
-func Test_Table_GranuleCompaction(t *testing.T) {
-	c, table := basicTable(t)
-	defer c.Close()
-	table.config.rowGroupSize = 2 // override row group size setting
-
-	samples := dynparquet.Samples{{
-		Labels: []dynparquet.Label{
-			{Name: "label1", Value: "value1"},
-			{Name: "label2", Value: "value2"},
-		},
-		Stacktrace: []uuid.UUID{
-			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
-			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
-		},
-		Timestamp: 1,
-		Value:     1,
-	}, {
-		Labels: []dynparquet.Label{
-			{Name: "label1", Value: "value1"},
-			{Name: "label2", Value: "value2"},
-			{Name: "label3", Value: "value3"},
-		},
-		Stacktrace: []uuid.UUID{
-			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
-			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
-		},
-		Timestamp: 2,
-		Value:     2,
-	}, {
-		Labels: []dynparquet.Label{
-			{Name: "label1", Value: "value1"},
-			{Name: "label2", Value: "value2"},
-			{Name: "label4", Value: "value4"},
-		},
-		Stacktrace: []uuid.UUID{
-			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x1},
-			{0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x2},
-		},
-		Timestamp: 3,
-		Value:     3,
-	}}
-
-	ctx := context.Background()
-	tx, err := table.Write(ctx, samples[0], samples[1], samples[2])
-	require.NoError(t, err)
-
-	// Wait for the writes to have been marked as completed
-	table.db.Wait(tx)
-
-	// Expect only a single granule to have been created
-	require.Equal(t, 1, table.active.Index().Len())
-
-	// Force compaction on the granule
-	g := (table.active.Index().Min()).(*Granule)
-	table.active.wg.Add(1)
-	table.active.compact(g)
-
-	// Expect the granule to have been compacted into a single granule
-	require.Equal(t, 1, table.active.Index().Len())
-
-	// Expect the granule to contain a single part with multiple row groups
-	g = (table.active.Index().Min()).(*Granule)
-	parts := 0
-	g.parts.Iterate(func(p *Part) bool {
-		parts++
-		require.Equal(t, 2, p.Buf.NumRowGroups())
-		return true
-	})
-
-	// There should only be a single part that contains multiple row groups
-	require.Equal(t, 1, parts)
-}