TTL mem-cache

go.mod

@@ -1,3 +1,14 @@
 module github.com/ukashazia/memcache
 
 go 1.23.0
+
+require (
+	github.com/google/uuid v1.6.0
+	github.com/stretchr/testify v1.10.0
+)
+
+require (
+	github.com/davecgh/go-spew v1.1.1 // indirect
+	github.com/pmezard/go-difflib v1.0.0 // indirect
+	gopkg.in/yaml.v3 v3.0.1 // indirect
+)

go.sum

@@ -0,0 +1,12 @@
+github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
+github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/google/uuid v1.6.0 h1:NIvaJDMOsjHA8n1jAhLSgzrAzy1Hgr+hNrb57e+94F0=
+github.com/google/uuid v1.6.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo=
+github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
+github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA=
+github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
+gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
+gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

ttl/ttl.go

@@ -4,6 +4,8 @@ import (
 	"context"
 	"sync"
 	"time"
+
+	"github.com/google/uuid"
 )
 
 type TTL struct {
@@ -23,9 +25,11 @@ type Item struct {
 }
 
 type shard struct {
-	id    uint64
-	data  map[string]*Item
-	mutex sync.RWMutex
+	id       uint64
+	uuid     uuid.UUID
+	data     map[string]*Item
+	termFunc context.CancelFunc
+	mutex    sync.RWMutex
 }
 
 type shards map[uint64]*shard
@@ -39,30 +43,34 @@ type shardLookupTable struct {
 func (ttl *TTL) Init() error {
 
 	newShard := shard{}
-	ttl.shardLookupTable = shardLookupTable{shards: shards{}}
+	ttl.shardLookupTable = shardLookupTable{shards: make(shards)}
 
 	ttl.newShard(&newShard)
 
 	return nil
 }
 
 func (ttl *TTL) Put(item *Item) (uint64, error) {
+
+	ttl.shardLookupTable.mutex.RLock()
 	shardId := ttl.shardLookupTable.currentShardId
 	currentShard := ttl.shardLookupTable.shards[shardId]
+	ttl.shardLookupTable.mutex.RUnlock()
 
 	if item.TTL.Nanoseconds() > 0 {
 		item.expirationTime = time.Now().Add(item.TTL)
 	} else {
 		item.expirationTime = time.Now().Add(ttl.DefaultTTL)
 	}
 
+	currentShard.mutex.Lock()
 	if uint64(len(currentShard.data)) < ttl.ShardSize {
 
-		currentShard.mutex.Lock()
 		defer currentShard.mutex.Unlock()
 
 		currentShard.data[item.Key] = item
 	} else {
+		currentShard.mutex.Unlock()
 
 		newShard := shard{}
 		ttl.newShard(&newShard)
@@ -78,7 +86,17 @@ func (ttl *TTL) Put(item *Item) (uint64, error) {
 }
 
 func (ttl *TTL) Get(key string, shardId uint64) any {
-	data, exists := ttl.shardLookupTable.shards[shardId].data[key]
+	ttl.shardLookupTable.mutex.RLock()
+	shard, exists := ttl.shardLookupTable.shards[shardId]
+	ttl.shardLookupTable.mutex.RUnlock()
+
+	if !exists {
+		return nil
+	}
+
+	shard.mutex.RLock()
+	data, exists := shard.data[key]
+	shard.mutex.RUnlock()
 
 	if !exists {
 		return nil
@@ -95,7 +113,12 @@ func (ttl *TTL) Get(key string, shardId uint64) any {
 }
 
 func (ttl *TTL) Delete(key string, shardId uint64) {
-	shard := ttl.shardLookupTable.shards[shardId]
+	shard, exists := ttl.shardLookupTable.shards[shardId]
+
+	if !exists {
+		return
+	}
+
 	shard.mutex.Lock()
 	defer shard.mutex.Unlock()
 
@@ -108,13 +131,16 @@ func (ttl *TTL) newShard(shard *shard) {
 	defer ttl.shardLookupTable.mutex.Unlock()
 
 	newShardId := ttl.shardLookupTable.currentShardId + 1
+	ctx, cancel := context.WithCancel(context.Background())
+
 	shard.id = newShardId
 	shard.data = make(map[string]*Item)
+	shard.uuid = uuid.New()
+	shard.termFunc = cancel
 
 	ttl.shardLookupTable.shards[newShardId] = shard
 	ttl.shardLookupTable.currentShardId = newShardId
 
-	ctx := context.Background()
 	go shard.cleanup(ctx, ttl)
 }
 
@@ -129,24 +155,41 @@ func (shard *shard) cleanup(ctx context.Context, ttl *TTL) {
 		case <-ticker.C:
 			var expiredKeys []string
 
+			shard.mutex.Lock()
 			for k, v := range shard.data {
 				if v.expirationTime.Before(time.Now()) {
 					expiredKeys = append(expiredKeys, k)
 				}
 			}
 
 			if len(expiredKeys) > 0 {
-				shard.mutex.Lock()
 				for _, k := range expiredKeys {
 					delete(shard.data, k)
 				}
-				if len(shard.data) == 0 {
-					delete(ttl.shardLookupTable.shards, shard.id) // idk if deleting the shard which is references in its own cleanup goroutine would work
+			}
 
-					return
-				}
+			shardEmpty := len(shard.data) == 0
+
+			if shardEmpty {
 				shard.mutex.Unlock()
+				ttl.terminateShard(shard)
+				return
 			}
+			shard.mutex.Unlock()
 		}
 	}
 }
+
+func (ttl *TTL) terminateShard(shard *shard) {
+
+	ttl.shardLookupTable.mutex.Lock()
+	defer ttl.shardLookupTable.mutex.Unlock()
+
+	shard.mutex.Lock()
+	defer shard.mutex.Unlock()
+
+	if _, exists := ttl.shardLookupTable.shards[shard.id]; exists {
+		shard.termFunc()
+		delete(ttl.shardLookupTable.shards, shard.id)
+	}
+}

ttl/ttl_test.go

@@ -0,0 +1,97 @@
+package ttl_test
+
+import (
+	"strconv"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/ukashazia/memcache/ttl"
+)
+
+func TestPutAndGet(t *testing.T) {
+	cache := &ttl.TTL{
+		DefaultTTL:      1 * time.Second,
+		ShardSize:       2,
+		CleanupInterval: 500 * time.Millisecond,
+	}
+	cache.Init()
+
+	item := &ttl.Item{Key: "test", Value: "value", TTL: 1 * time.Second}
+	shardID, err := cache.Put(item)
+	assert.Nil(t, err)
+
+	val := cache.Get("test", shardID)
+	assert.Equal(t, "value", val)
+}
+
+func TestExpiration(t *testing.T) {
+	cache := &ttl.TTL{
+		DefaultTTL:      500 * time.Millisecond,
+		ShardSize:       2,
+		CleanupInterval: 250 * time.Millisecond,
+	}
+	cache.Init()
+
+	item := &ttl.Item{Key: "temp", Value: "to expire", TTL: 500 * time.Millisecond}
+	shardID, _ := cache.Put(item)
+
+	time.Sleep(600 * time.Millisecond)
+	val := cache.Get("temp", shardID)
+	assert.Nil(t, val)
+}
+
+func TestShardCreation(t *testing.T) {
+	cache := &ttl.TTL{
+		DefaultTTL:      1 * time.Second,
+		ShardSize:       1,
+		CleanupInterval: 500 * time.Millisecond,
+	}
+	cache.Init()
+
+	item1 := &ttl.Item{Key: "item1", Value: "data1", TTL: 1 * time.Second}
+	shard1, _ := cache.Put(item1)
+	item2 := &ttl.Item{Key: "item2", Value: "data2", TTL: 1 * time.Second}
+	shard2, _ := cache.Put(item2)
+
+	assert.NotEqual(t, shard1, shard2)
+}
+
+func TestConcurrentAccess(t *testing.T) {
+	cache := &ttl.TTL{
+		DefaultTTL:      1 * time.Second,
+		ShardSize:       10,
+		CleanupInterval: 500 * time.Millisecond,
+	}
+	cache.Init()
+
+	var wg sync.WaitGroup
+	n := 100
+	for i := 0; i < n; i++ {
+		wg.Add(1)
+		go func(i int) {
+			defer wg.Done()
+			item := &ttl.Item{Key: strconv.Itoa(i), Value: i, TTL: 1 * time.Second}
+			_, _ = cache.Put(item)
+		}(i)
+	}
+
+	wg.Wait()
+	assert.NotEmpty(t, cache.Get("1", 1))
+}
+
+func TestDelete(t *testing.T) {
+	cache := &ttl.TTL{
+		DefaultTTL:      1 * time.Second,
+		ShardSize:       2,
+		CleanupInterval: 500 * time.Millisecond,
+	}
+	cache.Init()
+
+	item := &ttl.Item{Key: "delete_me", Value: "gone", TTL: 1 * time.Second}
+	shardID, _ := cache.Put(item)
+	cache.Delete("delete_me", shardID)
+
+	assert.Nil(t, cache.Get("delete_me", shardID))
+}