summaryrefslogtreecommitdiff
path: root/bst
diff options
context:
space:
mode:
Diffstat (limited to 'bst')
-rw-r--r--bst/bst.go152
-rw-r--r--bst/bst_test.go325
-rw-r--r--bst/doc.go7
-rw-r--r--bst/interface.go13
4 files changed, 497 insertions, 0 deletions
diff --git a/bst/bst.go b/bst/bst.go
new file mode 100644
index 00000000..dab9346c
--- /dev/null
+++ b/bst/bst.go
@@ -0,0 +1,152 @@
+package bst
+
+// BST ...
+type BST struct {
+ // registered topic, not unique
+ topic string
+ // associated connections with the topic
+ uuids map[string]struct{}
+
+ // left and right subtrees
+ left *BST
+ right *BST
+}
+
+func NewBST() Storage {
+ return &BST{
+ uuids: make(map[string]struct{}, 10),
+ }
+}
+
+// Insert uuid to the topic
+func (b *BST) Insert(uuid string, topic string) {
+ curr := b
+
+ for {
+ if topic == curr.topic {
+ curr.uuids[uuid] = struct{}{}
+ return
+ }
+ // if topic less than curr topic
+ if topic < curr.topic {
+ if curr.left == nil {
+ curr.left = &BST{
+ topic: topic,
+ uuids: map[string]struct{}{uuid: {}},
+ }
+ return
+ }
+ // move forward
+ curr = curr.left
+ } else {
+ if curr.right == nil {
+ curr.right = &BST{
+ topic: topic,
+ uuids: map[string]struct{}{uuid: {}},
+ }
+ return
+ }
+
+ curr = curr.right
+ }
+ }
+}
+
+func (b *BST) Contains(topic string) bool {
+ curr := b
+ for curr != nil {
+ switch {
+ case topic < curr.topic:
+ curr = curr.left
+ case topic > curr.topic:
+ curr = curr.right
+ case topic == curr.topic:
+ return true
+ }
+ }
+
+ return false
+}
+
+func (b *BST) Get(topic string) map[string]struct{} {
+ curr := b
+ for curr != nil {
+ switch {
+ case topic < curr.topic:
+ curr = curr.left
+ case topic > curr.topic:
+ curr = curr.right
+ case topic == curr.topic:
+ return curr.uuids
+ }
+ }
+
+ return nil
+}
+
+func (b *BST) Remove(uuid string, topic string) {
+ b.removeHelper(uuid, topic, nil)
+}
+
+func (b *BST) removeHelper(uuid string, topic string, parent *BST) {
+ curr := b
+ for curr != nil {
+ if topic < curr.topic { //nolint:gocritic
+ parent = curr
+ curr = curr.left
+ } else if topic > curr.topic {
+ parent = curr
+ curr = curr.right
+ } else {
+ // if more than 1 topic - remove only topic, do not remove the whole vertex
+ if len(curr.uuids) > 1 {
+ if _, ok := curr.uuids[uuid]; ok {
+ delete(curr.uuids, uuid)
+ return
+ }
+ }
+
+ if curr.left != nil && curr.right != nil { //nolint:gocritic
+ curr.topic, curr.uuids = curr.right.traverseForMinString()
+ curr.right.removeHelper(curr.topic, uuid, curr)
+ } else if parent == nil {
+ if curr.left != nil { //nolint:gocritic
+ curr.topic = curr.left.topic
+ curr.uuids = curr.left.uuids
+
+ curr.right = curr.left.right
+ curr.left = curr.left.left
+ } else if curr.right != nil {
+ curr.topic = curr.right.topic
+ curr.uuids = curr.right.uuids
+
+ curr.left = curr.right.left
+ curr.right = curr.right.right
+ } else { //nolint:staticcheck
+ // single node tree
+ }
+ } else if parent.left == curr {
+ if curr.left != nil {
+ parent.left = curr.left
+ } else {
+ parent.left = curr.right
+ }
+ } else if parent.right == curr {
+ if curr.left != nil {
+ parent.right = curr.left
+ } else {
+ parent.right = curr.right
+ }
+ }
+ break
+ }
+ }
+}
+
+//go:inline
+func (b *BST) traverseForMinString() (string, map[string]struct{}) {
+ if b.left == nil {
+ return b.topic, b.uuids
+ }
+ return b.left.traverseForMinString()
+}
diff --git a/bst/bst_test.go b/bst/bst_test.go
new file mode 100644
index 00000000..2271508c
--- /dev/null
+++ b/bst/bst_test.go
@@ -0,0 +1,325 @@
+package bst
+
+import (
+ "math/rand"
+ "testing"
+ "time"
+
+ "github.com/google/uuid"
+ "github.com/stretchr/testify/assert"
+)
+
+const predifined = "chat-1-2"
+
+func TestNewBST(t *testing.T) {
+ // create a new bst
+ g := NewBST()
+
+ for i := 0; i < 100; i++ {
+ g.Insert(uuid.NewString(), "comments")
+ }
+
+ for i := 0; i < 100; i++ {
+ g.Insert(uuid.NewString(), "comments2")
+ }
+
+ for i := 0; i < 100; i++ {
+ g.Insert(uuid.NewString(), "comments3")
+ }
+
+ // should be 100
+ exist := g.Get("comments")
+ assert.Len(t, exist, 100)
+
+ // should be 100
+ exist2 := g.Get("comments2")
+ assert.Len(t, exist2, 100)
+
+ // should be 100
+ exist3 := g.Get("comments3")
+ assert.Len(t, exist3, 100)
+}
+
+func BenchmarkGraph(b *testing.B) {
+ g := NewBST()
+
+ for i := 0; i < 1000; i++ {
+ uid := uuid.New().String()
+ g.Insert(uuid.NewString(), uid)
+ }
+
+ g.Insert(uuid.NewString(), predifined)
+
+ b.ResetTimer()
+ b.ReportAllocs()
+
+ for i := 0; i < b.N; i++ {
+ exist := g.Get(predifined)
+ _ = exist
+ }
+}
+
+func BenchmarkBigSearch(b *testing.B) {
+ g1 := NewBST()
+ g2 := NewBST()
+ g3 := NewBST()
+
+ predefinedSlice := make([]string, 0, 1000)
+ for i := 0; i < 1000; i++ {
+ predefinedSlice = append(predefinedSlice, uuid.NewString())
+ }
+ if predefinedSlice == nil {
+ b.FailNow()
+ }
+
+ for i := 0; i < 1000; i++ {
+ g1.Insert(uuid.NewString(), uuid.NewString())
+ }
+ for i := 0; i < 1000; i++ {
+ g2.Insert(uuid.NewString(), uuid.NewString())
+ }
+ for i := 0; i < 1000; i++ {
+ g3.Insert(uuid.NewString(), uuid.NewString())
+ }
+
+ for i := 0; i < 333; i++ {
+ g1.Insert(uuid.NewString(), predefinedSlice[i])
+ }
+
+ for i := 0; i < 333; i++ {
+ g2.Insert(uuid.NewString(), predefinedSlice[333+i])
+ }
+
+ for i := 0; i < 333; i++ {
+ g3.Insert(uuid.NewString(), predefinedSlice[666+i])
+ }
+
+ b.ResetTimer()
+ b.ReportAllocs()
+
+ for i := 0; i < b.N; i++ {
+ for i := 0; i < 333; i++ {
+ exist := g1.Get(predefinedSlice[i])
+ _ = exist
+ }
+ }
+ for i := 0; i < b.N; i++ {
+ for i := 0; i < 333; i++ {
+ exist := g2.Get(predefinedSlice[333+i])
+ _ = exist
+ }
+ }
+ for i := 0; i < b.N; i++ {
+ for i := 0; i < 333; i++ {
+ exist := g3.Get(predefinedSlice[666+i])
+ _ = exist
+ }
+ }
+}
+
+func BenchmarkBigSearchWithRemoves(b *testing.B) {
+ g1 := NewBST()
+ g2 := NewBST()
+ g3 := NewBST()
+
+ predefinedSlice := make([]string, 0, 1000)
+ for i := 0; i < 1000; i++ {
+ predefinedSlice = append(predefinedSlice, uuid.NewString())
+ }
+ if predefinedSlice == nil {
+ b.FailNow()
+ }
+
+ for i := 0; i < 1000; i++ {
+ g1.Insert(uuid.NewString(), uuid.NewString())
+ }
+ for i := 0; i < 1000; i++ {
+ g2.Insert(uuid.NewString(), uuid.NewString())
+ }
+ for i := 0; i < 1000; i++ {
+ g3.Insert(uuid.NewString(), uuid.NewString())
+ }
+
+ for i := 0; i < 333; i++ {
+ g1.Insert(uuid.NewString(), predefinedSlice[i])
+ }
+
+ for i := 0; i < 333; i++ {
+ g2.Insert(uuid.NewString(), predefinedSlice[333+i])
+ }
+
+ for i := 0; i < 333; i++ {
+ g3.Insert(uuid.NewString(), predefinedSlice[666+i])
+ }
+
+ go func() {
+ tt := time.NewTicker(time.Millisecond)
+ for {
+ select {
+ case <-tt.C:
+ num := rand.Intn(333) //nolint:gosec
+ values := g1.Get(predefinedSlice[num])
+ for k := range values {
+ g1.Remove(k, predefinedSlice[num])
+ }
+ }
+ }
+ }()
+
+ b.ResetTimer()
+ b.ReportAllocs()
+
+ for i := 0; i < b.N; i++ {
+ for i := 0; i < 333; i++ {
+ exist := g1.Get(predefinedSlice[i])
+ _ = exist
+ }
+ }
+ for i := 0; i < b.N; i++ {
+ for i := 0; i < 333; i++ {
+ exist := g2.Get(predefinedSlice[333+i])
+ _ = exist
+ }
+ }
+ for i := 0; i < b.N; i++ {
+ for i := 0; i < 333; i++ {
+ exist := g3.Get(predefinedSlice[666+i])
+ _ = exist
+ }
+ }
+}
+
+func TestGraph(t *testing.T) {
+ g := NewBST()
+
+ for i := 0; i < 1000; i++ {
+ uid := uuid.New().String()
+ g.Insert(uuid.NewString(), uid)
+ }
+
+ g.Insert(uuid.NewString(), predifined)
+
+ exist := g.Get(predifined)
+ assert.NotNil(t, exist)
+ assert.Len(t, exist, 1)
+}
+
+func TestTreeConcurrentContains(t *testing.T) {
+ g := NewBST()
+
+ key1 := uuid.NewString()
+ key2 := uuid.NewString()
+ key3 := uuid.NewString()
+ key4 := uuid.NewString()
+ key5 := uuid.NewString()
+
+ g.Insert(key1, predifined)
+ g.Insert(key2, predifined)
+ g.Insert(key3, predifined)
+ g.Insert(key4, predifined)
+ g.Insert(key5, predifined)
+
+ for i := 0; i < 100; i++ {
+ go func() {
+ _ = g.Get(predifined)
+ }()
+
+ go func() {
+ _ = g.Get(predifined)
+ }()
+
+ go func() {
+ _ = g.Get(predifined)
+ }()
+
+ go func() {
+ _ = g.Get(predifined)
+ }()
+ }
+
+ time.Sleep(time.Second * 2)
+
+ exist := g.Get(predifined)
+ assert.NotNil(t, exist)
+ assert.Len(t, exist, 5)
+}
+
+func TestGraphRemove(t *testing.T) {
+ g := NewBST()
+
+ key1 := uuid.NewString()
+ key2 := uuid.NewString()
+ key3 := uuid.NewString()
+ key4 := uuid.NewString()
+ key5 := uuid.NewString()
+
+ g.Insert(key1, predifined)
+ g.Insert(key2, predifined)
+ g.Insert(key3, predifined)
+ g.Insert(key4, predifined)
+ g.Insert(key5, predifined)
+
+ exist := g.Get(predifined)
+ assert.NotNil(t, exist)
+ assert.Len(t, exist, 5)
+
+ g.Remove(key1, predifined)
+
+ exist = g.Get(predifined)
+ assert.NotNil(t, exist)
+ assert.Len(t, exist, 4)
+}
+
+func TestBigSearch(t *testing.T) {
+ g1 := NewBST()
+ g2 := NewBST()
+ g3 := NewBST()
+
+ predefinedSlice := make([]string, 0, 1000)
+ for i := 0; i < 1000; i++ {
+ predefinedSlice = append(predefinedSlice, uuid.NewString())
+ }
+ if predefinedSlice == nil {
+ t.FailNow()
+ }
+
+ for i := 0; i < 1000; i++ {
+ g1.Insert(uuid.NewString(), uuid.NewString())
+ }
+ for i := 0; i < 1000; i++ {
+ g2.Insert(uuid.NewString(), uuid.NewString())
+ }
+ for i := 0; i < 1000; i++ {
+ g3.Insert(uuid.NewString(), uuid.NewString())
+ }
+
+ for i := 0; i < 333; i++ {
+ g1.Insert(uuid.NewString(), predefinedSlice[i])
+ }
+
+ for i := 0; i < 333; i++ {
+ g2.Insert(uuid.NewString(), predefinedSlice[333+i])
+ }
+
+ for i := 0; i < 333; i++ {
+ g3.Insert(uuid.NewString(), predefinedSlice[666+i])
+ }
+
+ for i := 0; i < 333; i++ {
+ exist := g1.Get(predefinedSlice[i])
+ assert.NotNil(t, exist)
+ assert.Len(t, exist, 1)
+ }
+
+ for i := 0; i < 333; i++ {
+ exist := g2.Get(predefinedSlice[333+i])
+ assert.NotNil(t, exist)
+ assert.Len(t, exist, 1)
+ }
+
+ for i := 0; i < 333; i++ {
+ exist := g3.Get(predefinedSlice[666+i])
+ assert.NotNil(t, exist)
+ assert.Len(t, exist, 1)
+ }
+}
diff --git a/bst/doc.go b/bst/doc.go
new file mode 100644
index 00000000..abb7e6e9
--- /dev/null
+++ b/bst/doc.go
@@ -0,0 +1,7 @@
+package bst
+
+/*
+Binary search tree for the pubsub
+
+The vertex may have one or multiply topics associated with the single websocket connection UUID
+*/
diff --git a/bst/interface.go b/bst/interface.go
new file mode 100644
index 00000000..95b03e11
--- /dev/null
+++ b/bst/interface.go
@@ -0,0 +1,13 @@
+package bst
+
+// Storage is general in-memory BST storage implementation
+type Storage interface {
+ // Insert inserts to a vertex with topic ident connection uuid
+ Insert(uuid string, topic string)
+ // Remove removes uuid from topic, if the uuid is single for a topic, whole vertex will be removed
+ Remove(uuid, topic string)
+ // Get will return all connections associated with the topic
+ Get(topic string) map[string]struct{}
+ // Contains checks if the BST contains a topic
+ Contains(topic string) bool
+}
generated by cgit v1.2.3 (git 2.25.1) at 2025-01-23 13:19:32 +0000