move prefix handler to pkg.

internal/proxy/manager.go

@@ -16,6 +16,8 @@ import (
 	dockerClient "github.com/docker/docker/client"
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
+
+	"github.com/m1k1o/neko-rooms/pkg/prefix"
 )
 
 type entry struct {
@@ -41,7 +43,7 @@ type ProxyManagerCtx struct {
 
 	client       *dockerClient.Client
 	instanceName string
-	handlers     *prefixHandler[*entry]
+	handlers     prefix.Tree[*entry]
 }
 
 func New(client *dockerClient.Client, instanceName string, waitEnabled bool) *ProxyManagerCtx {
@@ -52,7 +54,7 @@ func New(client *dockerClient.Client, instanceName string, waitEnabled bool) *Pr
 		client:       client,
 		instanceName: instanceName,
 		waitEnabled:  waitEnabled,
-		handlers:     &prefixHandler[*entry]{},
+		handlers:     prefix.NewTree[*entry](),
 	}
 }
 
@@ -119,7 +121,7 @@ func (p *ProxyManagerCtx) Start() {
 				p.mu.Lock()
 				switch msg.Action {
 				case "create":
-					p.handlers.Set(path, &entry{
+					p.handlers.Insert(path, &entry{
 						id:      msg.ID,
 						running: false,
 					})
@@ -137,9 +139,9 @@ func (p *ProxyManagerCtx) Start() {
 						})
 					}
 
-					p.handlers.Set(path, e)
+					p.handlers.Insert(path, e)
 				case "stop":
-					p.handlers.Set(path, &entry{
+					p.handlers.Insert(path, &entry{
 						id:      msg.ID,
 						running: false,
 					})
@@ -172,7 +174,7 @@ func (p *ProxyManagerCtx) Refresh() error {
 		return err
 	}
 
-	p.handlers = &prefixHandler[*entry]{}
+	p.handlers = prefix.NewTree[*entry]()
 
 	for _, cont := range containers {
 		enabled, path, port, ok := p.parseLabels(cont.Labels)
@@ -198,7 +200,7 @@ func (p *ProxyManagerCtx) Refresh() error {
 			}
 		}
 
-		p.handlers.Set(path, entry)
+		p.handlers.Insert(path, entry)
 	}
 
 	return nil

pkg/prefix/doc.go

@@ -0,0 +1,6 @@
+/*
+Package prefix implements a tree data structure that can be used to store
+and retrieve values based on a path prefix. It can be used to match paths
+against a set of prefixes.
+*/
+package prefix

pkg/prefix/tree.go

@@ -1,22 +1,29 @@
-package proxy
+package prefix
 
 import (
 	"strings"
 )
 
-type prefixHandler[T any] struct {
+type Tree[T any] interface {
+	Insert(prefix string, value T)
+	Find(prefix string) (value T, ok bool)
+	Remove(prefix string)
+	Match(path string) (value T, prefix string, ok bool)
+}
+
+type tree[T any] struct {
 	Value    T
 	IsLeaf   bool
-	Children map[string]*prefixHandler[T]
+	Children map[string]*tree[T]
 }
 
-func NewPrefixHandler[T any]() *prefixHandler[T] {
-	return &prefixHandler[T]{}
+func NewTree[T any]() *tree[T] {
+	return &tree[T]{}
 }
 
-func (p *prefixHandler[T]) Set(prefix string, value T) {
+func (p *tree[T]) Insert(prefix string, value T) {
 	arr := strings.Split(prefix, "/")
-	len := len(arr)
+	l := len(arr)
 
 	for i, a := range arr {
 		if a == "" {
@@ -24,15 +31,15 @@ func (p *prefixHandler[T]) Set(prefix string, value T) {
 		}
 
 		if p.Children == nil {
-			p.Children = map[string]*prefixHandler[T]{}
+			p.Children = map[string]*tree[T]{}
 		}
 
 		dat, ok := p.Children[a]
 		if !ok {
-			dat = &prefixHandler[T]{}
+			dat = &tree[T]{}
 		}
 
-		if i == len-1 {
+		if i == l-1 {
 			dat.Value = value
 			dat.IsLeaf = true
 			dat.Children = nil
@@ -43,9 +50,9 @@ func (p *prefixHandler[T]) Set(prefix string, value T) {
 	}
 }
 
-func (p *prefixHandler[T]) Get(prefix string) (value T, ok bool) {
+func (p *tree[T]) Find(prefix string) (value T, ok bool) {
 	arr := strings.Split(prefix, "/")
-	len := len(arr)
+	l := len(arr)
 
 	for i, a := range arr {
 		if a == "" {
@@ -53,15 +60,15 @@ func (p *prefixHandler[T]) Get(prefix string) (value T, ok bool) {
 		}
 
 		if p.Children == nil {
-			p.Children = map[string]*prefixHandler[T]{}
+			p.Children = map[string]*tree[T]{}
 		}
 
 		dat, found := p.Children[a]
 		if !found {
-			dat = &prefixHandler[T]{}
+			dat = &tree[T]{}
 		}
 
-		if i == len-1 {
+		if i == l-1 {
 			value = dat.Value
 			ok = dat.IsLeaf
 			return
@@ -74,16 +81,17 @@ func (p *prefixHandler[T]) Get(prefix string) (value T, ok bool) {
 	return
 }
 
-func (p *prefixHandler[T]) Remove(prefix string) {
+func (p *tree[T]) Remove(prefix string) {
 	arr := strings.Split(prefix, "/")
-	len := len(arr)
+	l := len(arr)
 
+	ptrs := []*tree[T]{p}
 	for i, a := range arr {
 		if a == "" {
 			continue
 		}
 
-		if i == len-1 {
+		if i == l-1 {
 			delete(p.Children, a)
 		}
 
@@ -92,12 +100,23 @@ func (p *prefixHandler[T]) Remove(prefix string) {
 			break
 		}
 		p = dat
+		ptrs = append(ptrs, p)
 	}
 
-	// TODO: Remove all empty references.
+	// remove all empty references
+	rm := false
+	for i := len(ptrs) - 1; i >= 0; i-- {
+		if len(ptrs[i].Children) == 0 {
+			rm = true
+			continue
+		}
+		if rm {
+			ptrs[i].Children = nil
+		}
+	}
 }
 
-func (p *prefixHandler[T]) Match(path string) (value T, prefix string, ok bool) {
+func (p *tree[T]) Match(path string) (value T, prefix string, ok bool) {
 	arr := strings.Split(path, "/")
 	prefixArr := []string{}
 

pkg/prefix/tree_test.go

@@ -0,0 +1,143 @@
+package prefix
+
+import (
+	"fmt"
+	"testing"
+)
+
+func TestMatch(t *testing.T) {
+	// create a new prefix tree
+	tree := NewTree[string]()
+
+	// add some values to the tree
+	tree.Insert("/users/1", "user1")
+	tree.Insert("/users/2", "user2")
+
+	// test matching a path that exists in the tree
+	value, prefix, ok := tree.Match("/users/1")
+	if !ok || value != "user1" || prefix != "/users/1" {
+		t.Errorf("Match failed for /users/1")
+	}
+
+	// test matching a path that exists in the tree but has a trailing slash
+	value, prefix, ok = tree.Match("/users/2/")
+	if !ok || value != "user2" || prefix != "/users/2" {
+		t.Errorf("Match failed for /users/2/")
+	}
+
+	// test matching a path that does not exist in the tree
+	_, _, ok = tree.Match("/users/4")
+	if ok {
+		t.Errorf("Match should have failed for /users/4")
+	}
+
+	// test matching a path that partially exists in the tree
+	value, prefix, ok = tree.Match("/users/1/posts")
+	if !ok || value != "user1" || prefix != "/users/1" {
+		t.Errorf("Match failed for /users/1/posts: %v %v %v", value, prefix, ok)
+	}
+}
+
+func TestFind(t *testing.T) {
+	// create a new prefix tree
+	tree := NewTree[string]()
+
+	// add some values to the tree
+	tree.Insert("/users/1", "user1")
+	tree.Insert("/users/2", "user2")
+
+	// test finding a path that exists in the tree
+	value, ok := tree.Find("/users/1")
+	if !ok || value != "user1" {
+		t.Errorf("Find failed for /users/1")
+	}
+
+	// test finding a path that exists in the tree but has a trailing slash
+	_, ok = tree.Find("/users/2/")
+	if ok {
+		t.Errorf("Find should have failed for /users/2/")
+	}
+
+	// test finding a path that does not exist in the tree
+	_, ok = tree.Find("/users/4")
+	if ok {
+		t.Errorf("Find should have failed for /users/4")
+	}
+
+	// test finding a path that partially exists in the tree
+	_, ok = tree.Find("/users/1/posts")
+	if ok {
+		t.Errorf("Find should have failed for /users/1/posts")
+	}
+}
+
+func TestRemove(t *testing.T) {
+	// create a new prefix tree
+	tree := &tree[string]{}
+
+	// add some values to the tree
+	tree.Insert("/users/1", "user1")
+	tree.Insert("/users/2", "user2")
+
+	// test matching a path that exists in the tree
+	value, prefix, ok := tree.Match("/users/1")
+	if !ok || value != "user1" || prefix != "/users/1" {
+		t.Errorf("Match failed for /users/1")
+	}
+
+	// remove a path that exists in the tree
+	tree.Remove("/users/1")
+
+	// test matching a path that exists in the tree
+	_, _, ok = tree.Match("/users/1")
+	if ok {
+		t.Errorf("Match should have failed for /users/1")
+	}
+
+	// test finding a path that exists in the tree
+	_, ok = tree.Find("/users/1")
+	if ok {
+		t.Errorf("Find should have failed for /users/1")
+	}
+
+	// remove another path that exists in the tree
+	tree.Remove("/users/2")
+
+	// test matching a path that exists in the tree
+	_, _, ok = tree.Match("/users/2")
+	if ok {
+		t.Errorf("Match should have failed for /users/2")
+	}
+
+	// check that the tree is empty
+	if tree.Children != nil {
+		t.Errorf("Tree should be empty, but contains %d children", len(tree.Children))
+		for k, v := range tree.Children {
+			t.Errorf("%s: %v", k, v)
+		}
+	}
+}
+
+// Benchmark the match function
+func Benchmark(b *testing.B) {
+	// create a new prefix tree
+	tree := NewTree[string]()
+
+	// insert 1 - 10 levels
+	level := "/users"
+	for i := 1; i <= 10; i++ {
+		level += fmt.Sprintf("/%d", i)
+		tree.Insert(level, fmt.Sprintf("level_%d", i))
+	}
+
+	// run the benchmark
+	level = "/users"
+	for i := 0; i <= 10; i++ {
+		level += fmt.Sprintf("/%d", i)
+		b.Run(fmt.Sprintf("level_%d", i), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				tree.Match(level)
+			}
+		})
+	}
+}