m1.2 : Added TTL support

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)
 

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

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=

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)
 	}

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)
-}
+	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
+		}
+	}
+}
+