diff --git a/cmd/server/main.go b/cmd/server/main.go index 7cbe340..f527ea4 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -37,6 +37,9 @@ func main() { if err != nil { log.Fatal("unable to open wal file: ", err) } + + go store.StartExpire(ctx) + server := server.NewTCPServer(port, store, wal) log.Printf("Starting server on %s\n", port) diff --git a/go.mod b/go.mod index 86cb2a7..3f0b1e5 100644 --- a/go.mod +++ b/go.mod @@ -1,3 +1,7 @@ module com.github.SantanuKar43/simple-kv go 1.25.6 + +require github.com/hashicorp/go-immutable-radix/v2 v2.1.0 + +require github.com/hashicorp/golang-lru/v2 v2.0.0 // indirect diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..6cf9c05 --- /dev/null +++ b/go.sum @@ -0,0 +1,8 @@ +github.com/hashicorp/go-immutable-radix/v2 v2.1.0 h1:CUW5RYIcysz+D3B+l1mDeXrQ7fUvGGCwJfdASSzbrfo= +github.com/hashicorp/go-immutable-radix/v2 v2.1.0/go.mod h1:hgdqLXA4f6NIjRVisM1TJ9aOJVNRqKZj+xDGF6m7PBw= +github.com/hashicorp/go-uuid v1.0.3 h1:2gKiV6YVmrJ1i2CKKa9obLvRieoRGviZFL26PcT/Co8= +github.com/hashicorp/go-uuid v1.0.3/go.mod h1:6SBZvOh/SIDV7/2o3Jml5SYk/TvGqwFJ/bN7x4byOro= +github.com/hashicorp/golang-lru/v2 v2.0.0 h1:Lf+9eD8m5pncvHAOCQj49GSN6aQI8XGfI5OpXNkoWaA= +github.com/hashicorp/golang-lru/v2 v2.0.0/go.mod h1:QeFd9opnmA6QUJc5vARoKUSoFhyfM2/ZepoAG6RGpeM= +golang.org/x/exp v0.0.0-20221215174704-0915cd710c24 h1:6w3iSY8IIkp5OQtbYj8NeuKG1jS9d+kYaubXqsoOiQ8= +golang.org/x/exp v0.0.0-20221215174704-0915cd710c24/go.mod h1:CxIveKay+FTh1D0yPZemJVgC/95VzuuOLq5Qi4xnoYc= diff --git a/internal/handler/handler.go b/internal/handler/handler.go index a9f784c..25c9eb4 100644 --- a/internal/handler/handler.go +++ b/internal/handler/handler.go @@ -1,10 +1,10 @@ package handler - + import ( "fmt" "log" - "strings" - + "strconv" + "time" "com.github.SantanuKar43/simple-kv/internal/protocol" "com.github.SantanuKar43/simple-kv/internal/store" "com.github.SantanuKar43/simple-kv/internal/wal" @@ -17,7 +17,8 @@ func Handle(input string, store *store.Store, wal *wal.Wal) string { if len(cmd.Args) != 2 { return "ERR wrong number of arguments" } - if err := appendToWal(strings.Join([]string{cmd.Name, cmd.Args[0], cmd.Args[1]}, " "), wal); err != nil { + walEntry := fmt.Sprintf("%s %s %s", cmd.Name, cmd.Args[0], cmd.Args[1]) + if err := appendToWal(walEntry, wal); err != nil { return fmt.Sprintf("ERR unable to write to WAL: %s", err.Error()) } @@ -36,12 +37,57 @@ func Handle(input string, store *store.Store, wal *wal.Wal) string { if len(cmd.Args) != 1 { return "ERR wrong number of arguments" } - if err := appendToWal(strings.Join([]string{cmd.Name, cmd.Args[0]}, " "), wal); err != nil { + walEntry := fmt.Sprintf("%s %s", cmd.Name, cmd.Args[0]) + if err := appendToWal(walEntry, wal); err != nil { return fmt.Sprintf("ERR unable to write to WAL: %s", err.Error()) } store.Delete(cmd.Args[0]) return "OK" + case "EXPIRE": + if len(cmd.Args) != 2 { + return "ERR wrong number of arguments" + } + sec, err := strconv.Atoi(cmd.Args[1]) + if err != nil { + return "ERR invalid expire time" + } + + expiresAt := time.Now().Add(time.Duration(sec) * time.Second).Unix() + + walEntry := fmt.Sprintf("%s %s %s %d", cmd.Name, cmd.Args[0], cmd.Args[1], expiresAt) + if err := appendToWal(walEntry, wal); err != nil { + return fmt.Sprintf("ERR unable to write to WAL: %s", err.Error()) + } + + if store.Expire(cmd.Args[0], time.Duration(sec)*time.Second) { + return "(integer) 1" + } + return "(integer) 0" + case "SETEX": + if len(cmd.Args) != 3 { + return "ERR wrong number of arguments" + } + sec, err := strconv.Atoi(cmd.Args[1]) + if err != nil { + return "ERR invalid expire time" + } + + expiresAt := time.Now().Add(time.Duration(sec) * time.Second).Unix() + + walEntry := fmt.Sprintf("%s %s %s %d %s", cmd.Name, cmd.Args[0], cmd.Args[1], expiresAt, cmd.Args[2]) + + if err := appendToWal(walEntry, wal); err != nil { + return fmt.Sprintf("ERR unable to write to WAL: %s", err.Error()) + } + + store.SetWithTTL(cmd.Args[0], cmd.Args[2], time.Duration(sec)*time.Second) + return "OK" + case "TTL": + if len(cmd.Args) != 1 { + return "ERR wrong number of arguments" + } + return fmt.Sprintf("(integer) %d", store.TTL(cmd.Args[0])) default: log.Println("unknown command", input) return fmt.Sprintf("ERR unknown command %s", cmd.Name) @@ -68,6 +114,38 @@ func Replay(input string, store *store.Store) { store.Set(cmd.Args[0], cmd.Args[1]) case "DEL": store.Delete(cmd.Args[0]) + case "SETEX": + if len(cmd.Args) < 4 { + log.Printf("invalid number of arguments for SETEX: %d", len(cmd.Args)) + return + } + expiresAt, err := strconv.ParseInt(cmd.Args[2], 10, 64) + if err != nil { + log.Printf("invalid expiresAt for SETEX: %s", cmd.Args[2]) + return + } + if expiresAt <= time.Now().Unix() { + log.Printf("expiresAt for SETEX is in the past: %d", expiresAt) + return + } + remaining := time.Duration(expiresAt - time.Now().Unix()) + store.SetWithTTL(cmd.Args[0], cmd.Args[3], remaining*time.Second) + case "EXPIRE": + if len(cmd.Args) < 3 { + log.Printf("invalid number of arguments for EXPIRE: %d", len(cmd.Args)) + return + } + expiresAt, err := strconv.ParseInt(cmd.Args[2], 10, 64) + if err != nil { + log.Printf("invalid expiresAt for EXPIRE: %s", cmd.Args[2]) + return + } + if expiresAt <= time.Now().Unix() { + log.Printf("expiresAt for EXPIRE is in the past: %d", expiresAt) + return + } + remaining := time.Duration(expiresAt - time.Now().Unix()) + store.Expire(cmd.Args[0], remaining*time.Second) default: log.Printf("invalid command parsed %s\n", cmd.Name) } diff --git a/internal/store/store.go b/internal/store/store.go index ea5031f..e99b664 100644 --- a/internal/store/store.go +++ b/internal/store/store.go @@ -2,36 +2,182 @@ package store import ( "sync" + "time" + "context" + "github.com/hashicorp/go-immutable-radix/v2" ) +const sampleSize = 20; + +type entry struct { + value string + expireAt int64 +} + type Store struct{ - mu sync.RWMutex - data map[string]string + mu sync.RWMutex + tree *iradix.Tree[entry] + lastExpiryKey []byte } func NewStore() *Store { return &Store{ - mu: sync.RWMutex{}, - data: make(map[string]string), + tree: iradix.New[entry](), } } func (s *Store) Get(key string) (string, bool) { s.mu.RLock() - defer s.mu.RUnlock() - val, ok := s.data[key] - return val, ok + snapshot := s.tree + s.mu.RUnlock() + + value, ok := snapshot.Get([]byte(key)) + if !ok { + return "", false + } + if value.expireAt > 0 && value.expireAt < time.Now().Unix() { + return "", false + } + return value.value, ok } func (s *Store) Set(key string, val string) { s.mu.Lock() defer s.mu.Unlock() - s.data[key] = val + txn := s.tree.Txn() + txn.Insert([]byte(key), entry{value: val, expireAt: 0}) + s.tree = txn.Commit() } - func (s *Store) Delete(key string) { s.mu.Lock() defer s.mu.Unlock() - delete(s.data, key) -} \ No newline at end of file + txn := s.tree.Txn() + txn.Delete([]byte(key)) + s.tree = txn.Commit() + +} + +func (s *Store) SetWithTTL(key string, val string, ttl time.Duration) { + s.mu.Lock() + defer s.mu.Unlock() + txn := s.tree.Txn() + txn.Insert([]byte(key), entry{value: val, expireAt: time.Now().Add(ttl).Unix()}) + s.tree = txn.Commit() +} + +func (s *Store) Expire(key string, ttl time.Duration) bool { + s.mu.Lock() + defer s.mu.Unlock() + txn := s.tree.Txn() + value, ok := txn.Get([]byte(key)) + if !ok { + return false + } + value.expireAt = time.Now().Add(ttl).Unix() + txn.Insert([]byte(key), value) + s.tree = txn.Commit() + return true +} + +func (s *Store) TTL(key string) int64 { + s.mu.RLock() + snapshot := s.tree + s.mu.RUnlock() + value, ok := snapshot.Get([]byte(key)) + if !ok { + return -2 // key does not exist + } + if value.expireAt == 0 { + return -1 // key exists, no TTL set + } + remaining := value.expireAt - time.Now().Unix() + if remaining < 0 { + return -2 // expired but not yet cleaned up + } + return remaining +} + +func (s *Store) activeExpire() { + for { + s.mu.RLock() + snapshot := s.tree + lastKey := s.lastExpiryKey + s.mu.RUnlock() + + + var expiredKeys []string + sampled := 0 + expiredCount := 0 + var newLastkey []byte + started := (lastKey == nil) + + snapshot.Root().Walk(func(key []byte, value entry) bool { + // skip till lastkey + if !started { + if string(key) == string(lastKey) { + started = true + return false + } + return false + } + + // past lastkey start sampling + if value.expireAt > 0 && value.expireAt < time.Now().Unix() { + expiredKeys = append(expiredKeys, string(key)) + expiredCount++ + } + sampled++ + newLastkey = key + if sampled >= sampleSize { + return true + } + return false + }) + + // Update position for next scan + s.mu.Lock() + if sampled < sampleSize { + // Didn't sample enough keys, wrapped around to end + s.lastExpiryKey = nil + } else { + // Save position to resume from next time + s.lastExpiryKey = newLastkey + } + s.mu.Unlock() + + // Delete expired keys + if len(expiredKeys) > 0 { + s.mu.Lock() + txn := s.tree.Txn() + for _, key := range expiredKeys { + val, ok := txn.Get([]byte(key)) + if ok && val.expireAt > 0 && val.expireAt < time.Now().Unix() { + txn.Delete([]byte(key)) + } + } + s.tree = txn.Commit() + s.mu.Unlock() + } + + // If < 25% expired, stop. Otherwise loop again (adaptive) + if float64(expiredCount) <= float64(sampleSize)*0.25 { + break + } + } +} + +func (s *Store) StartExpire(ctx context.Context) { + ticker := time.NewTicker(1 * time.Second) + defer ticker.Stop() + + for { + select { + case <-ticker.C: + s.activeExpire() + case <-ctx.Done(): + return + } + } +} +