diff --git a/store/tree/tree.go b/store/tree/tree.go deleted file mode 100644 index b00c421..0000000 --- a/store/tree/tree.go +++ /dev/null @@ -1,156 +0,0 @@ -package tree - -import ( - "github.com/emirpasic/gods/trees/redblacktree" - "regexp" - "strings" - "sync" - "time" -) - -type ExpireAtSecs int64 - -// Tree is a a thread-safe data structure for tracking expirable items. It automatically expires old entries and keys. -// It does not garbage collect. Items are only expired when interacting with the data structure. -type Tree struct { - sync.Mutex - defaultExpireAfterSecs int64 - tree *redblacktree.Tree -} - -func NewTree(expireAfterSecs int64) *Tree { - return &Tree{ - defaultExpireAfterSecs: expireAfterSecs, - tree: redblacktree.NewWithStringComparator(), - } -} - -// Keys returns a list of all valid keys in the tree, and a sum of every key's valid entries. -// It is expensive because it will result in the entire tree being counted and expired where necessary. -func (n *Tree) Keys() ([]string, int) { - var outKeys []string - var outCount int - - n.Lock() - keysI := n.tree.Keys() - n.Unlock() - - if keysI == nil { - return outKeys, outCount - } - - var key string - var keyCount int - for _, iface := range keysI { - key = iface.(string) - keyCount = n.Count(key) - outCount += keyCount - if keyCount == 0 { - continue - } - outKeys = append(outKeys, key) - } - - return outKeys, outCount -} - -// Count will return the number of entries at `entryKey`. It has the side effect of cleaning up -// stale entries and entry keys. -func (n *Tree) Count(entryKey string) int { - n.Lock() - defer n.Unlock() - - datesSecs := n.getAndCleanupUnsafe(entryKey) - if datesSecs == nil { - return 0 - } - - if len(*datesSecs) == 0 { - n.tree.Remove(entryKey) - return 0 - } - - datesSecs = removeExpired(datesSecs) - count := len(*datesSecs) - n.tree.Put(entryKey, *datesSecs) - return count -} - -// KeyMatch crawls the subtree to return keys starting with the `keyPattern` string. -func (n *Tree) KeyMatch(keyPattern string) []string { - var out []string - var wg sync.WaitGroup - re, err := regexp.Compile(strings.ReplaceAll(keyPattern, "*", "(^|$|.+)")) - if err != nil { - return []string{err.Error()} - } - - wg.Add(1) - go func() { - iterator := n.tree.Iterator() - var k string - var kOk bool - existed := iterator.Next() - for existed { - k, kOk = iterator.Key().(string) - if !kOk { - break - } - existed = iterator.Next() - if re.MatchString(k) { - if n.Count(k) > 0 { - out = append(out, k) - } - } - } - wg.Done() - }() - wg.Wait() - - return out -} - -func (n *Tree) Put(entryKey string) { - n.Lock() - defer n.Unlock() - - datesSecs := n.getAndCleanupUnsafe(entryKey) - if datesSecs == nil { - datesSecs = &[]int64{} - } - datesSecs = removeExpired(datesSecs) - secs := time.Now().Unix() - nextDatesSecs := append(*datesSecs, secs+n.defaultExpireAfterSecs) - n.tree.Put(entryKey, nextDatesSecs) -} - -// getAndCleanupUnsafe does not lock the mutex, so it can be used inside a lock -func (n *Tree) getAndCleanupUnsafe(entryKey string) *[]int64 { - val, found := n.tree.Get(entryKey) - if !found { - return nil - } - dates := val.([]int64) - if len(dates) == 0 { - // cleanup empty entry - n.tree.Remove(entryKey) - return nil - } - return &dates // not extra copy -} - -func removeExpired(datesSecs *[]int64) *[]int64 { - if len(*datesSecs) == 0 { - return datesSecs - } - currentTime := time.Now().Unix() - var out []int64 - // TODO: these are already sorted, so we can discard earlier entries - for _, removeAt := range *datesSecs { - if removeAt > currentTime { - // KEEP - not expired - out = append(out, removeAt) - } - } - return &out -} diff --git a/store/tree/tree_test.go b/store/tree/tree_test.go deleted file mode 100644 index c3beb59..0000000 --- a/store/tree/tree_test.go +++ /dev/null @@ -1,150 +0,0 @@ -package tree - -import ( - "github.com/stretchr/testify/assert" - "math/rand" - "testing" - "time" -) - -func TestTree_removeExpired(t *testing.T) { - keep1 := time.Now().Unix() + 5 - keep2 := time.Now().Unix() + 200 - entries := []int64{ - time.Now().Unix() - 2, // expired - time.Now().Unix() - 60, // expired - keep1, // KEEP - keep2, // KEEP - time.Now().Unix() - 1, // expired - } - - result := removeExpired(&entries) - assert.Equal(t, 2, len(*result)) - assert.Equal(t, (*result)[0], keep1) - assert.Equal(t, (*result)[1], keep2) -} - -func TestTree_Count(t *testing.T) { - t.Run("count returns number of non-expired and removes expired values", func(t *testing.T) { - tr := NewTree(2) - keys, entryCount := tr.Keys() - assert.Equal(t, 0, len(keys)) - assert.Equal(t, 0, entryCount) - tr.Put("willy") - tr.Put("willy") - tr.Put("Uncle Brick") - tr.Put("pander") - tr.Put("pander") - tr.Put("pander") - - keys, entryCount = tr.Keys() - assert.Equal(t, 3, len(keys)) - assert.Equal(t, 6, entryCount) - - // called twice in a row to make sure Count() doesn't trigger deletion or something weird - assert.Equal(t, 2, tr.Count("willy")) - assert.Equal(t, 2, tr.Count("willy")) - assert.Equal(t, 1, tr.Count("Uncle Brick")) - assert.Equal(t, 1, tr.Count("Uncle Brick")) - assert.Equal(t, 3, tr.Count("pander")) - assert.Equal(t, 3, tr.Count("pander")) - - // wrong case check - assert.Equal(t, 0, tr.Count("uncle brick")) - // unknowns - assert.Equal(t, 0, tr.Count("p")) - assert.Equal(t, 0, tr.Count("733")) - - time.Sleep(2 * time.Second) - - // count should force the tree to purge expired values - assert.Equal(t, 0, tr.Count("willy")) - assert.Equal(t, 0, tr.Count("Uncle Brick")) - assert.Equal(t, 0, tr.Count("pander")) - - keys, entryCount = tr.Keys() - assert.Equal(t, 0, len(keys)) // keys should be expired, but calling this will also expire them - assert.Equal(t, 0, entryCount) - - // now check that Keys() expires old keys - tr.Put("billy") - tr.Put("billy") - assert.Equal(t, 2, tr.Count("billy")) - keys, entryCount = tr.Keys() - assert.Equal(t, 1, len(keys)) - assert.Equal(t, 2, entryCount) - time.Sleep(2 * time.Second) - keys, entryCount = tr.Keys() - assert.Equal(t, 0, len(keys), "Keys() should expire keys") - assert.Equal(t, 0, entryCount, "Keys() should expire entries") - }) -} - -func TestTree_KeyMatch(t *testing.T) { - t.Run("returns only matches for a keyPattern", func(t *testing.T) { - tr := NewTree(60) - // we only put one key, but multiple counts - tr.Put("asdf") - tr.Put("asdf") - tr.Put("asdf") - - tr.Put("a:sdf") - tr.Put("a") - tr.Put("a:") - - tr.Put("mn:blah") - tr.Put("na:blahblah") - tr.Put("bla") - - tr.Put("a:elvis:5") - tr.Put("b:elvis:8:elvis") - tr.Put("e:elvis") - tr.Put("c:elvis:1") - - // curveballs - tr.Put("") // doesn't work - tr.Put(".+") - tr.Put("nil") - - assert.Equal(t, 3, len(tr.KeyMatch("^a:*"))) - - m := tr.KeyMatch("*bla*") - assert.Equalf(t, 3, len(m), "%+v", m) - - assert.ElementsMatch(t, []string{"a:elvis:5", "b:elvis:8:elvis", "e:elvis", "c:elvis:1"}, tr.KeyMatch("*elvis*")) - - // everything - all := tr.KeyMatch("*") - assert.Equalf(t, 14, len(all), "%+v", all) - }) - t.Run("it does not crash when empty", func(t *testing.T) { - tr := NewTree(5) - - assert.Equal(t, 0, len(tr.KeyMatch("asdf"))) - }) - t.Run("it does not return the key when all its values are expired", func(t *testing.T) { - tr := NewTree(1) - tr.Put("a") - tr.tree.Put("a", []string{}) - - tr.Put("cdbe") - tr.Put("cd:aa") - - result := tr.KeyMatch("cd") - assert.ElementsMatch(t, []string{"cdbe", "cd:aa"}, result) - - }) - t.Run("works with a huge tree", func(t *testing.T) { - charset := "abcdefghijklmnopqrstuvwxyz" - - tr := NewTree(10) - for i := 0; i < 100_000; i++ { - tr.Put( - string(charset[rand.Intn(len(charset))]) + - string(charset[rand.Intn(len(charset))]) + - string(charset[rand.Intn(len(charset))])) - } - assert.NotEmpty(t, tr.KeyMatch("a")) - }) - -}