95 lines
2.6 KiB
Go
95 lines
2.6 KiB
Go
package common
|
|
|
|
import (
|
|
"context"
|
|
"oc-discovery/conf"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/google/uuid"
|
|
)
|
|
|
|
// SearchIdleTimeout returns the configured search idle timeout (default 5s).
|
|
func SearchIdleTimeout() time.Duration {
|
|
if t := conf.GetConfig().SearchTimeout; t > 0 {
|
|
return time.Duration(t) * time.Second
|
|
}
|
|
return 5 * time.Second
|
|
}
|
|
|
|
// searchEntry holds the lifecycle state for one active search.
|
|
type searchEntry struct {
|
|
cancel context.CancelFunc
|
|
timer *time.Timer
|
|
idleTimeout time.Duration
|
|
}
|
|
|
|
// SearchTracker tracks one active search per user (peer or resource).
|
|
// Each search is keyed by a composite "user:searchID" so that a replaced
|
|
// search's late-arriving results can be told apart from the current one.
|
|
//
|
|
// Typical usage:
|
|
//
|
|
// ctx, cancel := context.WithCancel(parent)
|
|
// key := tracker.Register(userKey, cancel, idleTimeout)
|
|
// defer tracker.Cancel(key)
|
|
// // ... on each result: tracker.ResetIdle(key) + tracker.IsActive(key)
|
|
type SearchTracker struct {
|
|
mu sync.Mutex
|
|
entries map[string]*searchEntry
|
|
}
|
|
|
|
func NewSearchTracker() *SearchTracker {
|
|
return &SearchTracker{entries: map[string]*searchEntry{}}
|
|
}
|
|
|
|
// Register starts a new search for baseUser, cancelling any previous one.
|
|
// Returns the composite key "baseUser:searchID" to be used as the search identifier.
|
|
func (t *SearchTracker) Register(baseUser string, cancel context.CancelFunc, idleTimeout time.Duration) string {
|
|
compositeKey := baseUser + ":" + uuid.New().String()
|
|
t.mu.Lock()
|
|
t.cancelByPrefix(baseUser)
|
|
e := &searchEntry{cancel: cancel, idleTimeout: idleTimeout}
|
|
e.timer = time.AfterFunc(idleTimeout, func() { t.Cancel(compositeKey) })
|
|
t.entries[compositeKey] = e
|
|
t.mu.Unlock()
|
|
return compositeKey
|
|
}
|
|
|
|
// Cancel cancels the search(es) matching user (bare user key or composite key).
|
|
func (t *SearchTracker) Cancel(user string) {
|
|
t.mu.Lock()
|
|
t.cancelByPrefix(user)
|
|
t.mu.Unlock()
|
|
}
|
|
|
|
// ResetIdle resets the idle timer for compositeKey after a response arrives.
|
|
func (t *SearchTracker) ResetIdle(compositeKey string) {
|
|
t.mu.Lock()
|
|
if e, ok := t.entries[compositeKey]; ok {
|
|
e.timer.Reset(e.idleTimeout)
|
|
}
|
|
t.mu.Unlock()
|
|
}
|
|
|
|
// IsActive returns true if compositeKey is still the current active search.
|
|
func (t *SearchTracker) IsActive(compositeKey string) bool {
|
|
t.mu.Lock()
|
|
_, ok := t.entries[compositeKey]
|
|
t.mu.Unlock()
|
|
return ok
|
|
}
|
|
|
|
// cancelByPrefix cancels all entries whose key equals user or starts with "user:".
|
|
// Must be called with t.mu held.
|
|
func (t *SearchTracker) cancelByPrefix(user string) {
|
|
for k, e := range t.entries {
|
|
if k == user || strings.HasPrefix(k, user+":") {
|
|
e.timer.Stop()
|
|
e.cancel()
|
|
delete(t.entries, k)
|
|
}
|
|
}
|
|
}
|