Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
148 changes: 141 additions & 7 deletions maps/synclock_map.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package mapsutil
import (
"sync"
"sync/atomic"
"time"

"github.com/projectdiscovery/utils/errkit"
)
Expand All @@ -11,11 +12,24 @@ var (
ErrReadOnly = errkit.New("map is currently in read-only mode")
)

// EvictionEntry represents an entry with last access time
type EvictionEntry[K, V comparable] struct {
Key K
Value V
LastAccess time.Time
}

// SyncLock adds sync and lock capabilities to generic map
type SyncLockMap[K, V comparable] struct {
ReadOnly atomic.Bool
mu sync.RWMutex
Map Map[K, V]

// Eviction-related fields
inactivityDuration time.Duration
evictionMap map[K]*EvictionEntry[K, V]
evictionTicker *time.Ticker
stopEviction chan struct{}
}

type SyncLockMapOption[K, V comparable] func(slm *SyncLockMap[K, V])
Expand All @@ -26,6 +40,18 @@ func WithMap[K, V comparable](m Map[K, V]) SyncLockMapOption[K, V] {
}
}

// WithEviction enables inactivity-based eviction policy with the specified duration
func WithEviction[K, V comparable](inactivityDuration time.Duration) SyncLockMapOption[K, V] {
return func(slm *SyncLockMap[K, V]) {
slm.inactivityDuration = inactivityDuration
slm.evictionMap = make(map[K]*EvictionEntry[K, V])
slm.stopEviction = make(chan struct{})

// Start eviction goroutine
slm.startEvictionRoutine()
}
}

// NewSyncLockMap creates a new SyncLockMap.
// If an existing map is provided, it is used; otherwise, a new map is created.
func NewSyncLockMap[K, V comparable](options ...SyncLockMapOption[K, V]) *SyncLockMap[K, V] {
Expand All @@ -42,6 +68,52 @@ func NewSyncLockMap[K, V comparable](options ...SyncLockMapOption[K, V]) *SyncLo
return slm
}

// startEvictionRoutine starts the background eviction routine
func (s *SyncLockMap[K, V]) startEvictionRoutine() {
if s.inactivityDuration <= 0 {
return
}

// Check every 1/4 of the inactivity duration, but at least every 10ms
tickerInterval := s.inactivityDuration / 4
if tickerInterval < 10*time.Millisecond {
tickerInterval = 10 * time.Millisecond
}

s.evictionTicker = time.NewTicker(tickerInterval)

go func() {
for {
select {
case <-s.evictionTicker.C:
s.evictInactiveEntries()
case <-s.stopEviction:
return
}
}
}()
}

// evictInactiveEntries removes entries that have been inactive for too long
func (s *SyncLockMap[K, V]) evictInactiveEntries() {
s.mu.Lock()
defer s.mu.Unlock()

now := time.Now()
var keysToDelete []K

for k, entry := range s.evictionMap {
if now.Sub(entry.LastAccess) >= s.inactivityDuration {
keysToDelete = append(keysToDelete, k)
}
}

for _, k := range keysToDelete {
delete(s.Map, k)
delete(s.evictionMap, k)
}
}

// Lock the current map to read-only mode
func (s *SyncLockMap[K, V]) Lock() {
s.ReadOnly.Store(true)
Expand All @@ -61,26 +133,57 @@ func (s *SyncLockMap[K, V]) Set(k K, v V) error {
s.mu.Lock()
defer s.mu.Unlock()

s.Map[k] = v
now := time.Now()

// If eviction is enabled, handle eviction logic
if s.inactivityDuration > 0 {
// Update or create eviction entry
if entry, exists := s.evictionMap[k]; exists {
// Update existing entry
entry.Value = v
entry.LastAccess = now
} else {
// Create new entry
s.evictionMap[k] = &EvictionEntry[K, V]{
Key: k,
Value: v,
LastAccess: now,
}
}
}

s.Map[k] = v
return nil
}

// Get an item with syncronous access
func (s *SyncLockMap[K, V]) Get(k K) (V, bool) {
s.mu.RLock()
defer s.mu.RUnlock()

v, ok := s.Map[k]
s.mu.RUnlock()

// If eviction is enabled and key exists, update last access time
if s.inactivityDuration > 0 && ok {
s.mu.Lock()
if entry, exists := s.evictionMap[k]; exists {
entry.LastAccess = time.Now()
}
s.mu.Unlock()
}

return v, ok
}

// Get an item with syncronous access
// Delete an item with syncronous access
func (s *SyncLockMap[K, V]) Delete(k K) {
s.mu.Lock()
defer s.mu.Unlock()

// If eviction is enabled, clean up eviction tracking
if s.inactivityDuration > 0 {
delete(s.evictionMap, k)
}

delete(s.Map, k)
}

Expand All @@ -103,14 +206,45 @@ func (s *SyncLockMap[K, V]) Clone() *SyncLockMap[K, V] {
defer s.mu.Unlock()

smap := &SyncLockMap[K, V]{
ReadOnly: atomic.Bool{},
mu: sync.RWMutex{},
Map: s.Map.Clone(),
ReadOnly: atomic.Bool{},
mu: sync.RWMutex{},
Map: s.Map.Clone(),
inactivityDuration: s.inactivityDuration,
}
smap.ReadOnly.Store(s.ReadOnly.Load())

// If eviction is enabled, reinitialize eviction structures
if s.inactivityDuration > 0 {
smap.evictionMap = make(map[K]*EvictionEntry[K, V])
smap.stopEviction = make(chan struct{})

// Copy eviction entries with current time
now := time.Now()
for k, entry := range s.evictionMap {
smap.evictionMap[k] = &EvictionEntry[K, V]{
Key: k,
Value: entry.Value,
LastAccess: now,
}
}

// Start eviction routine for the cloned map
smap.startEvictionRoutine()
}

return smap
}

// StopEviction stops the background eviction routine
func (s *SyncLockMap[K, V]) StopEviction() {
if s.evictionTicker != nil {
s.evictionTicker.Stop()
}
if s.stopEviction != nil {
close(s.stopEviction)
}
}

// Has checks if the current map has the provided key
func (s *SyncLockMap[K, V]) Has(key K) bool {
s.mu.RLock()
Expand Down
141 changes: 140 additions & 1 deletion maps/synclock_map_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package mapsutil
import (
"errors"
"testing"
"time"

"github.com/stretchr/testify/require"
"golang.org/x/exp/maps"
Expand All @@ -19,7 +20,7 @@ func TestSyncLockMap(t *testing.T) {
}

t.Run("Test NewSyncLockMap with map ", func(t *testing.T) {
m := NewSyncLockMap[string, string](WithMap(Map[string, string]{
m := NewSyncLockMap(WithMap(Map[string, string]{
"key1": "value1",
"key2": "value2",
}))
Expand Down Expand Up @@ -173,3 +174,141 @@ func TestSyncLockMapWithConcurrency(t *testing.T) {
}
})
}

func TestSyncLockMapWithEviction(t *testing.T) {
t.Run("Test WithEviction option", func(t *testing.T) {
m := NewSyncLockMap(WithEviction[string, string](100 * time.Millisecond))
defer m.StopEviction()

require.NotNil(t, m.evictionMap, "eviction map should be initialized")
require.Equal(t, 100*time.Millisecond, m.inactivityDuration, "inactivity duration should be set")
})

t.Run("Test eviction after inactivity", func(t *testing.T) {
m := NewSyncLockMap(WithEviction[string, string](50 * time.Millisecond))
defer m.StopEviction()

// Add some items
err := m.Set("key1", "value1")
require.NoError(t, err)
err = m.Set("key2", "value2")
require.NoError(t, err)

// Verify items exist
require.True(t, m.Has("key1"))
require.True(t, m.Has("key2"))

// Wait for eviction (wait longer to ensure eviction happens)
time.Sleep(200 * time.Millisecond)

// Items should be evicted
require.False(t, m.Has("key1"))
require.False(t, m.Has("key2"))
})

t.Run("Test access resets eviction timer", func(t *testing.T) {
m := NewSyncLockMap(WithEviction[string, string](100 * time.Millisecond))
defer m.StopEviction()

// Add an item
err := m.Set("key1", "value1")
require.NoError(t, err)

// Access the item before eviction
time.Sleep(60 * time.Millisecond)
_, ok := m.Get("key1")
require.True(t, ok, "item should still exist after access")

// Wait a bit more but not enough for eviction
time.Sleep(60 * time.Millisecond)
_, ok = m.Get("key1")
require.True(t, ok, "item should still exist after recent access")

// Now wait for eviction
time.Sleep(150 * time.Millisecond)
_, ok = m.Get("key1")
require.False(t, ok, "item should be evicted after inactivity")
})

t.Run("Test Set updates access time", func(t *testing.T) {
m := NewSyncLockMap(WithEviction[string, string](100 * time.Millisecond))
defer m.StopEviction()

// Add an item
err := m.Set("key1", "value1")
require.NoError(t, err)

// Update the item before eviction
time.Sleep(60 * time.Millisecond)
err = m.Set("key1", "value1_updated")
require.NoError(t, err)

// Wait a bit more but not enough for eviction
time.Sleep(60 * time.Millisecond)
value, ok := m.Get("key1")
require.True(t, ok, "item should still exist after update")
require.Equal(t, "value1_updated", value, "value should be updated")

// Now wait for eviction
time.Sleep(150 * time.Millisecond)
_, ok = m.Get("key1")
require.False(t, ok, "item should be evicted after inactivity")
})

t.Run("Test Delete removes from eviction tracking", func(t *testing.T) {
m := NewSyncLockMap(WithEviction[string, string](50 * time.Millisecond))
defer m.StopEviction()

// Add an item
err := m.Set("key1", "value1")
require.NoError(t, err)

// Delete the item
m.Delete("key1")
require.False(t, m.Has("key1"))

// Wait and verify it's not in eviction map
time.Sleep(100 * time.Millisecond)
require.False(t, m.Has("key1"))
})

t.Run("Test Clone with eviction", func(t *testing.T) {
m := NewSyncLockMap(WithEviction[string, string](100 * time.Millisecond))
defer m.StopEviction()

// Add some items
err := m.Set("key1", "value1")
require.NoError(t, err)
err = m.Set("key2", "value2")
require.NoError(t, err)

// Clone the map
cloned := m.Clone()
defer cloned.StopEviction()

// Verify cloned map has the same items
require.True(t, cloned.Has("key1"))
require.True(t, cloned.Has("key2"))

// Verify cloned map has eviction enabled
require.Equal(t, m.inactivityDuration, cloned.inactivityDuration)
require.NotNil(t, cloned.evictionMap)
})

t.Run("Test StopEviction", func(t *testing.T) {
m := NewSyncLockMap(WithEviction[string, string](50 * time.Millisecond))

// Add an item
err := m.Set("key1", "value1")
require.NoError(t, err)

// Stop eviction
m.StopEviction()

// Wait for what would normally be eviction time
time.Sleep(100 * time.Millisecond)

// Item should still exist since eviction is stopped
require.True(t, m.Has("key1"))
})
}
4 changes: 3 additions & 1 deletion scripts/docupdater/docupdater.go
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,9 @@ func fetchDoc(url string) (string, error) {
if err != nil {
return "", err
}
defer resp.Body.Close()
defer func() {
_ = resp.Body.Close()
}()

if resp.StatusCode != http.StatusOK {
return "", fmt.Errorf("failed to fetch file: %s", resp.Status)
Expand Down
Loading