1 files changed, 32 insertions, 3 deletions
diff --git a/plugins/websockets/storage/storage.go b/plugins/websockets/storage/storage.go
index 34f53cfd..0557dee5 100644
--- a/plugins/websockets/storage/storage.go
+++ b/plugins/websockets/storage/storage.go
@@ -17,7 +17,7 @@ func NewStorage() *Storage {
 	}
 }
 
-func (s *Storage) Store(connID string, topics []string) {
+func (s *Storage) InsertMany(connID string, topics []string) {
 	s.Lock()
 	defer s.Unlock()
 
@@ -26,7 +26,14 @@ func (s *Storage) Store(connID string, topics []string) {
 	}
 }
 
-func (s *Storage) Remove(connID string, topics []string) {
+func (s *Storage) Insert(connID string, topic string) {
+	s.Lock()
+	defer s.Unlock()
+
+	s.BST.Insert(connID, topic)
+}
+
+func (s *Storage) RemoveMany(connID string, topics []string) {
 	s.Lock()
 	defer s.Unlock()
 
@@ -35,7 +42,29 @@ func (s *Storage) Remove(connID string, topics []string) {
 	}
 }
 
-func (s *Storage) Get(topics []string, res map[string]struct{}) {
+func (s *Storage) Remove(connID string, topic string) {
+	s.Lock()
+	defer s.Unlock()
+
+	s.BST.Remove(connID, topic)
+}
+
+// Thread safe get
+func (s *Storage) GetByPtrTS(topics []string, res map[string]struct{}) {
+	s.Lock()
+	defer s.Unlock()
+
+	for i := 0; i < len(topics); i++ {
+		d := s.BST.Get(topics[i])
+		if len(d) > 0 {
+			for ii := range d {
+				res[ii] = struct{}{}
+			}
+		}
+	}
+}
+
+func (s *Storage) GetByPtr(topics []string, res map[string]struct{}) {
 	s.RLock()
 	defer s.RUnlock()