Repository: numaproj/numaflow
Pull request #390 (2022-11-29 06:34:49)

chore: adding DeleteStore method to store provider interface
Signed-off-by: Yashash H L <yashashhl25@gmail.com>

- removing GC from the Store interface
- adding DeleteStore to StoreProvider
- tests for StoreProvider


Patch:
diff --git a/pkg/pbq/pbq.go b/pkg/pbq/pbq.go
index 94b65c29..11b7f351 100644
--- a/pkg/pbq/pbq.go
+++ b/pkg/pbq/pbq.go
@@ -87,13 +87,8 @@ func (p *PBQ) ReadCh() <-chan *isb.ReadMessage {
 // GC cleans up the PBQ and also the store associated with it. GC is invoked after the Reader (ProcessAndForward) has
 // finished forwarding the output to ISB.
 func (p *PBQ) GC() error {
-	err := p.store.GC()
-	if err != nil {
-		return err
-	}
 	p.store = nil
-	p.manager.deregister(p.PartitionID)
-	return nil
+	return p.manager.deregister(p.PartitionID)
 }
 
 // replayRecordsFromStore replays store messages when replay flag is set during start up time. It replays by reading from
diff --git a/pkg/pbq/pbqmanager.go b/pkg/pbq/pbqmanager.go
index 210de411..84bc67ec 100644
--- a/pkg/pbq/pbqmanager.go
+++ b/pkg/pbq/pbqmanager.go
@@ -195,11 +195,12 @@ func (m *Manager) register(partitionID partition.ID, p *PBQ) {
 }
 
 // deregister is intended to be used by PBQ to deregister itself after GC is called.
-func (m *Manager) deregister(partitionID partition.ID) {
+// it will also delete the store using the store provider
+func (m *Manager) deregister(partitionID partition.ID) error {
 	m.Lock()
 	defer m.Unlock()
-
 	delete(m.pbqMap, partitionID.String())
+	return m.storeProvider.DeleteStore(partitionID)
 }
 
 func (m *Manager) getPBQs() []*PBQ {
diff --git a/pkg/pbq/store/interfaces.go b/pkg/pbq/store/interfaces.go
index e51e874e..912782c7 100644
--- a/pkg/pbq/store/interfaces.go
+++ b/pkg/pbq/store/interfaces.go
@@ -32,8 +32,6 @@ type Store interface {
 	Write(msg *isb.ReadMessage) error
 	// Close closes store
 	Close() error
-	// GC does garbage collection and deletes all the messages that are persisted
-	GC() error
 }
 
 // StoreProvider defines the functions for store implementation
@@ -42,4 +40,6 @@ type StoreProvider interface {
 	CreateStore(context.Context, partition.ID) (Store, error)
 	// DiscoverPartitions discovers all the managed partitions.
 	DiscoverPartitions(context.Context) ([]partition.ID, error)
+	// DeleteStore deletes the store
+	DeleteStore(partition.ID) error
 }
diff --git a/pkg/pbq/store/memory/store.go b/pkg/pbq/store/memory/store.go
index b28e4ca5..5128d254 100644
--- a/pkg/pbq/store/memory/store.go
+++ b/pkg/pbq/store/memory/store.go
@@ -74,15 +74,6 @@ func (m *memoryStore) Close() error {
 	return nil
 }
 
-// GC does garbage collection
-// for in-memory implementation we set the storage to nil, so that it will
-// ready for GC
-func (m *memoryStore) GC() error {
-	m.storage = nil
-	m.writePos = -1
-	return nil
-}
-
 // isEmpty check if there are any records persisted in store
 func (m *memoryStore) isEmpty() bool {
 	// is empty should return true when the store is created and no messages are written
diff --git a/pkg/pbq/store/memory/stores.go b/pkg/pbq/store/memory/stores.go
index 4b4f61ef..c9d30a24 100644
--- a/pkg/pbq/store/memory/stores.go
+++ b/pkg/pbq/store/memory/stores.go
@@ -18,6 +18,8 @@ package memory
 
 import (
 	"context"
+	"errors"
+	"sync"
 
 	"github.com/numaproj/numaflow/pkg/isb"
 	"github.com/numaproj/numaflow/pkg/pbq/partition"
@@ -28,22 +30,16 @@ import (
 type memoryStores struct {
 	storeSize    int64
 	discoverFunc func(ctx context.Context) ([]partition.ID, error)
-	partitions   map[partition.ID]store.Store
+	partitions   map[partition.ID]*memoryStore
+	sync.RWMutex
 }
 
 func NewMemoryStores(opts ...Option) store.StoreProvider {
 	s := &memoryStores{
 		storeSize:  100,
-		partitions: make(map[partition.ID]store.Store),
-	}
-	// default discover function
-	s.discoverFunc = func(ctx context.Context) ([]partition.ID, error) {
-		partitionsIds := make([]partition.ID, 0)
-		for key := range s.partitions {
-			partitionsIds = append(partitionsIds, key)
-		}
-		return partitionsIds, nil
+		partitions: make(map[partition.ID]*memoryStore),
 	}
+
 	for _, o := range opts {
 		o(s)
 	}
@@ -51,6 +47,8 @@ func NewMemoryStores(opts ...Option) store.StoreProvider {
 }
 
 func (ms *memoryStores) CreateStore(ctx context.Context, partitionID partition.ID) (store.Store, error) {
+	ms.Lock()
+	defer ms.Unlock()
 	if memStore, ok := ms.partitions[partitionID]; ok {
 		return memStore, nil
 	}
@@ -64,10 +62,30 @@ func (ms *memoryStores) CreateStore(ctx context.Context, partitionID partition.I
 		partitionID: partitionID,
 	}
 	ms.partitions[partitionID] = memStore
-
 	return memStore, nil
 }
 
 func (ms *memoryStores) DiscoverPartitions(ctx context.Context) ([]partition.ID, error) {
+	if ms.discoverFunc == nil {
+		partitionsIds := make([]partition.ID, 0)
+		for key := range ms.partitions {
+			partitionsIds = append(partitionsIds, key)
+		}
+		return partitionsIds, nil
+	}
 	return ms.discoverFunc(ctx)
 }
+
+func (ms *memoryStores) DeleteStore(partitionID partition.ID) error {
+	ms.Lock()
+	defer ms.Unlock()
+	memStore, ok := ms.partitions[partitionID]
+	if !ok {
+		return errors.New("store not found")
+	}
+
+	memStore.storage = nil
+	memStore.writePos = -1
+	delete(ms.partitions, partitionID)
+	return nil
+}
diff --git a/pkg/pbq/store/noop/store.go b/pkg/pbq/store/noop/store.go
index 96a971c8..3f97bdf9 100644
--- a/pkg/pbq/store/noop/store.go
+++ b/pkg/pbq/store/noop/store.go
@@ -42,7 +42,3 @@ func (p *PBQNoOpStore) Write(msg *isb.ReadMessage) error {
 func (p *PBQNoOpStore) Close() error {
 	return nil
 }
-
-func (p *PBQNoOpStore) GC() error {
-	return nil
-}
diff --git a/pkg/pbq/store/noop/stores.go b/pkg/pbq/store/noop/stores.go
index 0f89c138..9e392718 100644
--- a/pkg/pbq/store/noop/stores.go
+++ b/pkg/pbq/store/noop/stores.go
@@ -37,3 +37,7 @@ func (ns *noopStores) CreateStore(ctx context.Context, partitionID partition.ID)
 func (ns *noopStores) DiscoverPartitions(ctx context.Context) ([]partition.ID, error) {
 	return []partition.ID{}, nil
 }
+
+func (ns *noopStores) DeleteStore(partitionID partition.ID) error {
+	return nil
+}
diff --git a/pkg/pbq/store/wal/metrics.go b/pkg/pbq/store/wal/metrics.go
index 8e00bdc3..5b2482ad 100644
--- a/pkg/pbq/store/wal/metrics.go
+++ b/pkg/pbq/store/wal/metrics.go
@@ -66,13 +66,6 @@ var entryWriteTime = promauto.NewHistogramVec(prometheus.HistogramOpts{
 	Buckets:   prometheus.ExponentialBucketsRange(1, 60, 5),
 }, []string{labelPartitionKey})
 
-var lifespan = promauto.NewHistogramVec(prometheus.HistogramOpts{
-	Subsystem: "pbq_wal",
-	Name:      "wal_lifespan",
-	Help:      "Lifespan of a pbq wal (1 to 20 minutes)",
-	Buckets:   prometheus.ExponentialBucketsRange(1, 20, 5),
-}, []string{labelPartitionKey})
-
 var walErrors = promauto.NewCounterVec(prometheus.CounterOpts{
 	Subsystem: "pbq_wal",
 	Name:      "wal_errors",
diff --git a/pkg/pbq/store/wal/segment.go b/pkg/pbq/store/wal/segment.go
index 335f871c..19ccac71 100644
--- a/pkg/pbq/store/wal/segment.go
+++ b/pkg/pbq/store/wal/segment.go
@@ -293,30 +293,6 @@ func (w *WAL) Close() (err error) {
 	return nil
 }
 
-// GC cleans up the WAL Segment.
-func (w *WAL) GC() (err error) {
-	defer func() {
-		if err != nil {
-			walErrors.With(map[string]string{"kind": "gc"}).Inc()
-		}
-	}()
-	start := time.Now()
-	if !w.closed {
-		err := w.Close()
-		if err != nil {
-			return err
-		}
-	}
-	err = os.Remove(w.fp.Name())
-
-	if err == nil {
-		garbageCollectingTime.With(map[string]string{labelPartitionKey: w.partitionID.Key}).Observe(float64(time.Since(start).Microseconds()))
-		lifespan.With(map[string]string{labelPartitionKey: w.partitionID.Key}).Observe(time.Since(w.createTime).Minutes())
-		activeFilesCount.With(map[string]string{}).Dec()
-	}
-	return err
-}
-
 func getSegmentFilePath(id *partition.ID, dir string) string {
 	filename := fmt.Sprintf("%s_%d.%d.%s", SegmentPrefix, id.Start.Unix(), id.End.Unix(), id.Key)
 	return filepath.Join(dir, filename)
diff --git a/pkg/pbq/store/wal/stores.go b/pkg/pbq/store/wal/stores.go
index 2436b91e..d9e4b995 100644
--- a/pkg/pbq/store/wal/stores.go
+++ b/pkg/pbq/store/wal/stores.go
@@ -77,7 +77,7 @@ func (ws *walStores) openOrCreateWAL(id *partition.ID) (*WAL, error) {
 	if os.IsNotExist(err) {
 		// here we are explicitly giving O_WRONLY because we will not be using this to read. Our read is only during
 		// boot up.
-		fp, err = os.OpenFile(getSegmentFilePath(id, ws.storePath), os.O_WRONLY|os.O_CREATE, 0644)
+		fp, err = os.OpenFile(filePath, os.O_WRONLY|os.O_CREATE, 0644)
 		if err != nil {
 			return nil, err
 		}
@@ -157,3 +157,29 @@ func (ws *walStores) DiscoverPartitions(ctx context.Context) ([]partition.ID, er
 
 	return partitions, nil
 }
+
+func (ws *walStores) DeleteStore(partitionID partition.ID) error {
+	var err error
+	defer func() {
+		if err != nil {
+			walErrors.With(map[string]string{"kind": "gc"}).Inc()
+		}
+	}()
+
+	filePath := getSegmentFilePath(&partitionID, ws.storePath)
+	_, err = os.Stat(filePath)
+
+	if err != nil {
+		return err
+	}
+
+	start := time.Now()
+	// an open file can also be deleted
+	err = os.Remove(filePath)
+
+	if err == nil {
+		garbageCollectingTime.With(map[string]string{labelPartitionKey: partitionID.Key}).Observe(float64(time.Since(start).Microseconds()))
+		activeFilesCount.With(map[string]string{}).Dec()
+	}
+	return err
+}


Test patch:
diff --git a/pkg/pbq/pbqmanager_test.go b/pkg/pbq/pbqmanager_test.go
index a453e340..234fc605 100644
--- a/pkg/pbq/pbqmanager_test.go
+++ b/pkg/pbq/pbqmanager_test.go
@@ -35,7 +35,6 @@ import (
 
 func TestManager_ListPartitions(t *testing.T) {
 	size := int64(100)
-	t.Skip()
 
 	ctx := context.Background()
 	pbqManager, err := NewManager(ctx, memory.NewMemoryStores(memory.WithStoreSize(size)),
diff --git a/pkg/pbq/store/memory/stores_test.go b/pkg/pbq/store/memory/stores_test.go
new file mode 100644
index 00000000..6d8ae399
--- /dev/null
+++ b/pkg/pbq/store/memory/stores_test.go
@@ -0,0 +1,56 @@
+package memory
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestMemoryStores(t *testing.T) {
+	var err error
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+	partitionIds := []partition.ID{
+		{
+			Start: time.Unix(60, 0),
+			End:   time.Unix(120, 0),
+			Key:   "test-1",
+		},
+		{
+			Start: time.Unix(120, 0),
+			End:   time.Unix(180, 0),
+			Key:   "test-2",
+		},
+		{
+			Start: time.Unix(180, 0),
+			End:   time.Unix(240, 0),
+			Key:   "test-3",
+		},
+	}
+	storeProvider := NewMemoryStores(WithStoreSize(100))
+
+	for _, partitionID := range partitionIds {
+		_, err := storeProvider.CreateStore(ctx, partitionID)
+		assert.NoError(t, err)
+	}
+
+	var discoveredPartitions []partition.ID
+	discoveredPartitions, err = storeProvider.DiscoverPartitions(ctx)
+	assert.NoError(t, err)
+
+	assert.Len(t, discoveredPartitions, len(partitionIds))
+
+	for _, partitionID := range partitionIds {
+		err = storeProvider.DeleteStore(partitionID)
+		assert.NoError(t, err)
+	}
+
+	discoveredPartitions, err = storeProvider.DiscoverPartitions(ctx)
+	assert.NoError(t, err)
+
+	assert.Len(t, discoveredPartitions, 0)
+}
diff --git a/pkg/pbq/store/wal/stores_test.go b/pkg/pbq/store/wal/stores_test.go
new file mode 100644
index 00000000..ccfc88f1
--- /dev/null
+++ b/pkg/pbq/store/wal/stores_test.go
@@ -0,0 +1,57 @@
+package wal
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/numaproj/numaflow/pkg/pbq/partition"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestWalStores(t *testing.T) {
+	var err error
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+	defer cancel()
+	partitionIds := []partition.ID{
+		{
+			Start: time.Unix(60, 0),
+			End:   time.Unix(120, 0),
+			Key:   "test-1",
+		},
+		{
+			Start: time.Unix(120, 0),
+			End:   time.Unix(180, 0),
+			Key:   "test-2",
+		},
+		{
+			Start: time.Unix(180, 0),
+			End:   time.Unix(240, 0),
+			Key:   "test-3",
+		},
+	}
+
+	tmp := t.TempDir()
+	storeProvider := NewWALStores(WithStorePath(tmp))
+
+	for _, partitionID := range partitionIds {
+		_, err = storeProvider.CreateStore(ctx, partitionID)
+		assert.NoError(t, err)
+	}
+
+	var discoveredPartitions []partition.ID
+	discoveredPartitions, err = storeProvider.DiscoverPartitions(ctx)
+	assert.NoError(t, err)
+
+	assert.Len(t, discoveredPartitions, len(partitionIds))
+
+	for _, partitionID := range partitionIds {
+		err = storeProvider.DeleteStore(partitionID)
+		assert.NoError(t, err)
+	}
+
+	discoveredPartitions, err = storeProvider.DiscoverPartitions(ctx)
+	assert.NoError(t, err)
+
+	assert.Len(t, discoveredPartitions, 0)
+}