- Add Clear method to the storages

Signed-off-by: Valery Piashchynski <[email protected]>

8 files changed, 275 insertions, 4 deletions

diff --git a/plugins/kv/drivers/boltdb/driver.go b/plugins/kv/drivers/boltdb/driver.go
index 4b675271..0b411e5e 100644
--- a/plugins/kv/drivers/boltdb/driver.go
+++ b/plugins/kv/drivers/boltdb/driver.go
@@ -373,6 +373,31 @@ func (d *Driver) TTL(keys ...string) (map[string]string, error) {
 	return m, nil
 }
 
+func (d *Driver) Clear() error {
+	err := d.DB.Update(func(tx *bolt.Tx) error {
+		err := tx.DeleteBucket(d.bucket)
+		if err != nil {
+			d.log.Error("boltdb delete bucket", "error", err)
+			return err
+		}
+
+		_, err = tx.CreateBucket(d.bucket)
+		if err != nil {
+			d.log.Error("boltdb create bucket", "error", err)
+			return err
+		}
+
+		return nil
+	})
+
+	if err != nil {
+		d.log.Error("clear transaction failed", "error", err)
+		return err
+	}
+
+	return nil
+}
+
 // ========================= PRIVATE =================================
 
 func (d *Driver) startGCLoop() { //nolint:gocognit
diff --git a/plugins/kv/drivers/memcached/driver.go b/plugins/kv/drivers/memcached/driver.go
index a2787d72..14e7c078 100644
--- a/plugins/kv/drivers/memcached/driver.go
+++ b/plugins/kv/drivers/memcached/driver.go
@@ -237,3 +237,13 @@ func (d *Driver) Delete(keys ...string) error {
 	}
 	return nil
 }
+
+func (d *Driver) Clear() error {
+	err := d.client.DeleteAll()
+	if err != nil {
+		d.log.Error("flush_all operation failed", "error", err)
+		return err
+	}
+
+	return nil
+}
diff --git a/plugins/kv/interface.go b/plugins/kv/interface.go
index ffdbbe62..5736a6a7 100644
--- a/plugins/kv/interface.go
+++ b/plugins/kv/interface.go
@@ -22,9 +22,12 @@ type Storage interface {
 	MExpire(items ...*kvv1.Item) error
 
 	// TTL return the rest time to live for provided keys
-	// Not supported for the memcached and boltdb
+	// Not supported for the memcached
 	TTL(keys ...string) (map[string]string, error)
 
+	// Clear clean the entire storage
+	Clear() error
+
 	// Delete one or multiple keys.
 	Delete(keys ...string) error
 }
diff --git a/plugins/kv/rpc.go b/plugins/kv/rpc.go
index af763600..3f7ba97c 100644
--- a/plugins/kv/rpc.go
+++ b/plugins/kv/rpc.go
@@ -161,3 +161,19 @@ func (r *rpc) Delete(in *kvv1.Request, _ *kvv1.Response) error {
 
 	return errors.E(op, errors.Errorf("no such storage: %s", in.GetStorage()))
 }
+
+// Clear clean the storage
+func (r *rpc) Clear(in *kvv1.Request, _ *kvv1.Response) error {
+	const op = errors.Op("rcp_delete")
+
+	if st, exists := r.storages[in.GetStorage()]; exists {
+		err := st.Clear()
+		if err != nil {
+			return errors.E(op, err)
+		}
+
+		return nil
+	}
+
+	return errors.E(op, errors.Errorf("no such storage: %s", in.GetStorage()))
+}
diff --git a/plugins/memory/kv.go b/plugins/memory/kv.go
index 1cf031d1..83f6630e 100644
--- a/plugins/memory/kv.go
+++ b/plugins/memory/kv.go
@@ -13,7 +13,8 @@ import (
 )
 
 type Driver struct {
-	heap sync.Map
+	clearMu sync.RWMutex
+	heap    sync.Map
 	// stop is used to stop keys GC and close boltdb connection
 	stop chan struct{}
 	log  logger.Logger
@@ -203,6 +204,14 @@ func (s *Driver) Delete(keys ...string) error {
 	return nil
 }
 
+func (s *Driver) Clear() error {
+	s.clearMu.Lock()
+	s.heap = sync.Map{}
+	s.clearMu.Unlock()
+
+	return nil
+}
+
 // ================================== PRIVATE ======================================
 
 func (s *Driver) gc() {
@@ -213,6 +222,9 @@ func (s *Driver) gc() {
 			ticker.Stop()
 			return
 		case now := <-ticker.C:
+			// mutes needed to clear the map
+			s.clearMu.Lock()
+
 			// check every second
 			s.heap.Range(func(key, value interface{}) bool {
 				v := value.(*kvv1.Item)
@@ -231,6 +243,8 @@ func (s *Driver) gc() {
 				}
 				return true
 			})
+
+			s.clearMu.Unlock()
 		}
 	}
 }
diff --git a/plugins/redis/kv.go b/plugins/redis/kv.go
index 320b7443..2e4b9bfd 100644
--- a/plugins/redis/kv.go
+++ b/plugins/redis/kv.go
@@ -240,3 +240,12 @@ func (d *Driver) TTL(keys ...string) (map[string]string, error) {
 	}
 	return m, nil
 }
+
+func (d *Driver) Clear() error {
+	fdb := d.universalClient.FlushDB(context.Background())
+	if fdb.Err() != nil {
+		return fdb.Err()
+	}
+
+	return nil
+}
diff --git a/tests/plugins/broadcast/broadcast_plugin_test.go b/tests/plugins/broadcast/broadcast_plugin_test.go
index 2cd4b451..0ec813f3 100644
--- a/tests/plugins/broadcast/broadcast_plugin_test.go
+++ b/tests/plugins/broadcast/broadcast_plugin_test.go
@@ -205,7 +205,8 @@ func TestBroadcastSameSubscriber(t *testing.T) {
 		cfg,
 		&broadcast.Plugin{},
 		&rpcPlugin.Plugin{},
-		mockLogger,
+		&logger.ZapLogger{},
+		// mockLogger,
 		&server.Plugin{},
 		&redis.Plugin{},
 		&websockets.Plugin{},
@@ -314,7 +315,8 @@ func TestBroadcastSameSubscriberGlobal(t *testing.T) {
 		cfg,
 		&broadcast.Plugin{},
 		&rpcPlugin.Plugin{},
-		mockLogger,
+		&logger.ZapLogger{},
+		// mockLogger,
 		&server.Plugin{},
 		&redis.Plugin{},
 		&websockets.Plugin{},
diff --git a/tests/plugins/kv/storage_plugin_test.go b/tests/plugins/kv/storage_plugin_test.go
index 1e466e06..6169fad5 100644
--- a/tests/plugins/kv/storage_plugin_test.go
+++ b/tests/plugins/kv/storage_plugin_test.go
@@ -575,6 +575,54 @@ func testRPCMethods(t *testing.T) {
 	err = client.Call("kv.Has", keysDel, ret)
 	assert.NoError(t, err)
 	assert.Len(t, ret.GetItems(), 0)
+
+	dataClear := &payload.Request{
+		Storage: "boltdb-rr",
+		Items: []*payload.Item{
+			{
+				Key:   "a",
+				Value: []byte("aa"),
+			},
+			{
+				Key:   "b",
+				Value: []byte("bb"),
+			},
+			{
+				Key:   "c",
+				Value: []byte("cc"),
+			},
+			{
+				Key:   "d",
+				Value: []byte("dd"),
+			},
+			{
+				Key:   "e",
+				Value: []byte("ee"),
+			},
+		},
+	}
+
+	clear := &payload.Request{Storage: "boltdb-rr"}
+
+	ret = &payload.Response{}
+	// Register 3 keys with values
+	err = client.Call("kv.Set", dataClear, ret)
+	assert.NoError(t, err)
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Has", dataClear, ret)
+	assert.NoError(t, err)
+	assert.Len(t, ret.GetItems(), 5) // should be 5
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Clear", clear, ret)
+	assert.NoError(t, err)
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Has", dataClear, ret)
+	assert.NoError(t, err)
+	assert.Len(t, ret.GetItems(), 0) // should be 5
+
 }
 
 func TestMemcached(t *testing.T) {
@@ -790,6 +838,55 @@ func testRPCMethodsMemcached(t *testing.T) {
 	err = client.Call("kv.Has", keysDel, ret)
 	assert.NoError(t, err)
 	assert.Len(t, ret.GetItems(), 0)
+
+	dataClear := &payload.Request{
+		Storage: "memcached-rr",
+		Items: []*payload.Item{
+			{
+				Key:   "a",
+				Value: []byte("aa"),
+			},
+			{
+				Key:   "b",
+				Value: []byte("bb"),
+			},
+			{
+				Key:   "c",
+				Value: []byte("cc"),
+			},
+			{
+				Key:   "d",
+				Value: []byte("dd"),
+			},
+			{
+				Key:   "e",
+				Value: []byte("ee"),
+			},
+		},
+	}
+
+	clear := &payload.Request{Storage: "memcached-rr"}
+
+	ret = &payload.Response{}
+	// Register 3 keys with values
+	err = client.Call("kv.Set", dataClear, ret)
+	assert.NoError(t, err)
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Has", dataClear, ret)
+	assert.NoError(t, err)
+	assert.Len(t, ret.GetItems(), 5) // should be 5
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Clear", clear, ret)
+	assert.NoError(t, err)
+
+	time.Sleep(time.Second * 2)
+	ret = &payload.Response{}
+	err = client.Call("kv.Has", dataClear, ret)
+	assert.NoError(t, err)
+	assert.Len(t, ret.GetItems(), 0) // should be 5
+
 }
 
 func TestInMemory(t *testing.T) {
@@ -1004,6 +1101,54 @@ func testRPCMethodsInMemory(t *testing.T) {
 	err = client.Call("kv.Has", keysDel, ret)
 	assert.NoError(t, err)
 	assert.Len(t, ret.GetItems(), 0)
+
+	dataClear := &payload.Request{
+		Storage: "memory-rr",
+		Items: []*payload.Item{
+			{
+				Key:   "a",
+				Value: []byte("aa"),
+			},
+			{
+				Key:   "b",
+				Value: []byte("bb"),
+			},
+			{
+				Key:   "c",
+				Value: []byte("cc"),
+			},
+			{
+				Key:   "d",
+				Value: []byte("dd"),
+			},
+			{
+				Key:   "e",
+				Value: []byte("ee"),
+			},
+		},
+	}
+
+	clear := &payload.Request{Storage: "memory-rr"}
+
+	ret = &payload.Response{}
+	// Register 3 keys with values
+	err = client.Call("kv.Set", dataClear, ret)
+	assert.NoError(t, err)
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Has", dataClear, ret)
+	assert.NoError(t, err)
+	assert.Len(t, ret.GetItems(), 5) // should be 5
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Clear", clear, ret)
+	assert.NoError(t, err)
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Has", dataClear, ret)
+	assert.NoError(t, err)
+	assert.Len(t, ret.GetItems(), 0) // should be 5
+
 }
 
 func TestRedis(t *testing.T) {
@@ -1354,4 +1499,51 @@ func testRPCMethodsRedis(t *testing.T) {
 	err = client.Call("kv.Has", keysDel, ret)
 	assert.NoError(t, err)
 	assert.Len(t, ret.GetItems(), 0)
+
+	dataClear := &payload.Request{
+		Storage: "redis-rr",
+		Items: []*payload.Item{
+			{
+				Key:   "a",
+				Value: []byte("aa"),
+			},
+			{
+				Key:   "b",
+				Value: []byte("bb"),
+			},
+			{
+				Key:   "c",
+				Value: []byte("cc"),
+			},
+			{
+				Key:   "d",
+				Value: []byte("dd"),
+			},
+			{
+				Key:   "e",
+				Value: []byte("ee"),
+			},
+		},
+	}
+
+	clear := &payload.Request{Storage: "redis-rr"}
+
+	ret = &payload.Response{}
+	// Register 3 keys with values
+	err = client.Call("kv.Set", dataClear, ret)
+	assert.NoError(t, err)
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Has", dataClear, ret)
+	assert.NoError(t, err)
+	assert.Len(t, ret.GetItems(), 5) // should be 5
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Clear", clear, ret)
+	assert.NoError(t, err)
+
+	ret = &payload.Response{}
+	err = client.Call("kv.Has", dataClear, ret)
+	assert.NoError(t, err)
+	assert.Len(t, ret.GetItems(), 0) // should be 5
 }