diff --git a/internal/proxy/manager.go b/internal/proxy/manager.go index 1ea2c81..2196d0a 100644 --- a/internal/proxy/manager.go +++ b/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 diff --git a/pkg/prefix/doc.go b/pkg/prefix/doc.go new file mode 100644 index 0000000..e1298f2 --- /dev/null +++ b/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 diff --git a/internal/proxy/prefix.go b/pkg/prefix/tree.go similarity index 54% rename from internal/proxy/prefix.go rename to pkg/prefix/tree.go index 5d895ee..5466b16 100644 --- a/internal/proxy/prefix.go +++ b/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{} diff --git a/pkg/prefix/tree_test.go b/pkg/prefix/tree_test.go new file mode 100644 index 0000000..a333103 --- /dev/null +++ b/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) + } + }) + } +}