initial commit

internals/traceIdTtlMap/traceIdTtlMap.go

@@ -0,0 +1,141 @@
+package traceIdTtlMap
+
+import (
+	"sync"
+	"time"
+)
+
+type TTLMap struct {
+	m        map[string]int64
+	mu       sync.RWMutex
+	maxTtl   int64
+	stopCh   chan struct{}
+	stopOnce sync.Once
+}
+
+/**
+ * Creates a new TTLMap with the given maximum TTL.
+ * The TTLMap will automatically remove expired trace ids from the map, every second.
+ * Map is thread-safe.
+ */
+func New(initSize int, maxTtl int) (m *TTLMap) {
+	m = &TTLMap{
+		m:      make(map[string]int64, initSize),
+		maxTtl: int64(maxTtl),
+		stopCh: make(chan struct{}),
+	}
+
+	go func() {
+		ticker := time.NewTicker(time.Second)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ticker.C:
+				var expiredKeys []string
+				now := time.Now().Unix()
+
+				m.mu.RLock()
+				for key, value := range m.m {
+					if value < now {
+						expiredKeys = append(expiredKeys, key)
+					}
+				}
+				m.mu.RUnlock()
+
+				for _, key := range expiredKeys {
+					m.mu.Lock()
+					delete(m.m, key)
+					m.mu.Unlock()
+				}
+			case <-m.stopCh:
+				return
+			}
+		}
+	}()
+
+	return
+}
+
+/**
+ * Stops all go routines.
+ * Should be called when the TTLMap is no longer needed.
+ */
+func (m *TTLMap) Stop() {
+	m.stopOnce.Do(func() {
+		close(m.stopCh)
+	})
+}
+
+/**
+ * Adds a trace id to the map.
+ * When providing the same trace id twice, the second invocation will be ignored.
+ */
+func (m *TTLMap) Add(key string) {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+
+	_, ok := m.m[key]
+	if ok {
+		return
+	}
+
+	m.m[key] = time.Now().Unix() + m.maxTtl
+}
+
+/**
+ * Checks if a trace id exists in the map.
+ * Returns true if the trace id exists and has not expired.
+ * Returns false if the trace id does not exist or has expired.
+ * Removes the trace id from the map if it has expired.
+ */
+func (m *TTLMap) Exists(key string) bool {
+	m.mu.RLock()
+	value, ok := m.m[key]
+	m.mu.RUnlock()
+
+	if ok {
+		if value < time.Now().Unix() {
+			m.mu.Lock()
+			delete(m.m, key)
+			m.mu.Unlock()
+			return false
+		}
+	}
+
+	return ok
+}
+
+/**
+ * Inserts a new entry into the map.
+ * Only used for testing.
+ */
+func (m *TTLMap) insertEntry(key string, value int64) {
+	m.mu.Lock()
+	m.m[key] = value
+	m.mu.Unlock()
+}
+
+/**
+ * Gets an entry from the map.
+ * Only used for testing.
+ */
+func (m *TTLMap) getEntry(key string) (int64, bool) {
+	m.mu.RLock()
+	value, ok := m.m[key]
+	m.mu.RUnlock()
+
+	return value, ok
+}
+
+/**
+ * Gets the size of the map.
+ * Only used for testing.
+ */
+func (m *TTLMap) getSize() int {
+	m.mu.RLock()
+	size := len(m.m)
+	m.mu.RUnlock()
+
+	return size
+}

internals/traceIdTtlMap/traceIdTtlMap_benchmark_test.go

@@ -0,0 +1,43 @@
+package traceIdTtlMap
+
+import (
+	"strconv"
+	"testing"
+)
+
+func BenchmarkTTLMap_AddExists(b *testing.B) {
+	m := New(1000, 10)
+	defer m.Stop()
+
+	b.Run("Add", func(b *testing.B) {
+		for i := 0; b.Loop(); i++ {
+			m.Add(strconv.Itoa(i))
+		}
+	})
+
+	for i := range 1000 {
+		m.Add(strconv.Itoa(i))
+	}
+
+	b.Run("Exists", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			key := strconv.Itoa(i % 1000)
+			m.Exists(key)
+		}
+	})
+}
+
+func BenchmarkTTLMap_Concurrent(b *testing.B) {
+	m := New(1000, 10)
+	defer m.Stop()
+
+	b.RunParallel(func(pb *testing.PB) {
+		i := 0
+		for pb.Next() {
+			key := strconv.Itoa(i % 1000)
+			m.Add(key)
+			m.Exists(key)
+			i++
+		}
+	})
+}

internals/traceIdTtlMap/traceIdTtlMap_test.go

@@ -0,0 +1,113 @@
+package traceIdTtlMap
+
+import (
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestNew(t *testing.T) {
+	m := New(10, 10)
+	defer m.Stop()
+
+	assert.Equal(t, 0, m.getSize())
+}
+
+func TestNew_Cleanup(t *testing.T) {
+	m := New(10, 10)
+	defer m.Stop()
+
+	m.insertEntry("4bf92f3577b34da6a3ce929d0e0e4736", time.Now().Unix()-10)
+	assert.Equal(t, 1, m.getSize())
+
+	// Inserted entry should be deleted after >1 second
+	time.Sleep(time.Second * 2)
+	assert.Equal(t, 0, m.getSize())
+}
+
+func TestAdd(t *testing.T) {
+	m := New(10, 10)
+	defer m.Stop()
+
+	m.Add("4bf92f3577b34da6a3ce929d0e0e4736")
+	// Intentionally adding the same trace id twice, should not add it again
+	m.Add("4bf92f3577b34da6a3ce929d0e0e4736")
+	m.Add("d0240fe9f68b48e687d25c185d4c17c5")
+
+	assert.Equal(t, 2, m.getSize())
+
+	_, ok := m.getEntry("4bf92f3577b34da6a3ce929d0e0e4736")
+	assert.True(t, ok)
+	_, ok = m.getEntry("d0240fe9f68b48e687d25c185d4c17c5")
+	assert.True(t, ok)
+}
+
+func TestAdd_ResetTtl(t *testing.T) {
+	m := New(10, 10)
+	defer m.Stop()
+
+	m.Add("4bf92f3577b34da6a3ce929d0e0e4736")
+	insertTime, _ := m.getEntry("4bf92f3577b34da6a3ce929d0e0e4736")
+
+	time.Sleep(time.Second)
+
+	m.Add("4bf92f3577b34da6a3ce929d0e0e4736")
+	updatedTime, _ := m.getEntry("4bf92f3577b34da6a3ce929d0e0e4736")
+
+	// Delete time of the second entry should remain the same.
+	assert.Equal(t, updatedTime, insertTime)
+}
+
+func TestExists(t *testing.T) {
+	m := New(10, 10)
+	defer m.Stop()
+
+	m.insertEntry("4bf92f3577b34da6a3ce929d0e0e4736", time.Now().Unix()+10)
+
+	// Existing and valid trace
+	assert.True(t, m.Exists("4bf92f3577b34da6a3ce929d0e0e4736"))
+	// Non existing trace
+	assert.False(t, m.Exists("d0240fe9f68b48e687d25c185d4c17c5"))
+}
+
+func TestExists_ExpiredTrace(t *testing.T) {
+	m := New(10, 10)
+	defer m.Stop()
+
+	m.insertEntry("4bf92f3577b34da6a3ce929d0e0e4736", time.Now().Unix()-10)
+
+	// Existing and but invalid trace
+	assert.False(t, m.Exists("4bf92f3577b34da6a3ce929d0e0e4736"))
+}
+
+func TestAddExists_Concurrent(t *testing.T) {
+	m := New(10, 10)
+	defer m.Stop()
+
+	var wg sync.WaitGroup
+	keys := []string{"k1", "k2", "k3", "k4", "k5", "k6", "k7", "k8", "k9", "k10"}
+
+	for i := range 100 {
+		wg.Add(2)
+
+		go func(i int) {
+			defer wg.Done()
+			m.Add(keys[i%len(keys)])
+		}(i)
+
+		go func(i int) {
+			defer wg.Done()
+			_ = m.Exists(keys[i%len(keys)])
+		}(i)
+	}
+
+	wg.Wait()
+}
+
+func TestStop(t *testing.T) {
+	m := New(10, 10)
+	m.Stop()
+	m.Stop()
+}