From 5f9c7813df1fdc18dd455ea0381510796a4db071 Mon Sep 17 00:00:00 2001 From: juancwu Date: Wed, 29 Apr 2026 12:31:00 +0000 Subject: [PATCH] add simple memory cache server --- go.mod | 3 + server.go | 296 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 299 insertions(+) create mode 100644 go.mod create mode 100644 server.go diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..bb56949 --- /dev/null +++ b/go.mod @@ -0,0 +1,3 @@ +module git.juancwu.dev/juancwu/cubby + +go 1.26.2 diff --git a/server.go b/server.go new file mode 100644 index 0000000..e58acbf --- /dev/null +++ b/server.go @@ -0,0 +1,296 @@ +package main + +import ( + "bufio" + "flag" + "fmt" + "log" + "net" + "os" + "os/signal" + "strconv" + "strings" + "sync" + "syscall" + "time" +) + +type entry struct { + value string + expiresAt time.Time // zero means no expiry +} + +type Store struct { + mu sync.RWMutex + data map[string]entry +} + +func NewStore() *Store { + s := &Store{data: make(map[string]entry)} + go s.janitor() + return s +} + +// janitor sweeps expired keys once per second. +func (s *Store) janitor() { + t := time.NewTicker(time.Second) + defer t.Stop() + for range t.C { + now := time.Now() + s.mu.Lock() + for k, e := range s.data { + if !e.expiresAt.IsZero() && now.After(e.expiresAt) { + delete(s.data, k) + } + } + s.mu.Unlock() + } +} + +func (s *Store) Get(k string) (string, bool) { + s.mu.RLock() + e, ok := s.data[k] + s.mu.RUnlock() + if !ok { + return "", false + } + if !e.expiresAt.IsZero() && time.Now().After(e.expiresAt) { + s.mu.Lock() + delete(s.data, k) + s.mu.Unlock() + return "", false + } + return e.value, true +} + +func (s *Store) Set(k, v string, ttl time.Duration) { + var exp time.Time + if ttl > 0 { + exp = time.Now().Add(ttl) + } + s.mu.Lock() + s.data[k] = entry{value: v, expiresAt: exp} + s.mu.Unlock() +} + +func (s *Store) Del(k string) bool { + s.mu.Lock() + _, existed := s.data[k] + delete(s.data, k) + s.mu.Unlock() + return existed +} + +func (s *Store) Keys() []string { + now := time.Now() + s.mu.RLock() + defer s.mu.RUnlock() + out := make([]string, 0, len(s.data)) + for k, e := range s.data { + if e.expiresAt.IsZero() || now.Before(e.expiresAt) { + out = append(out, k) + } + } + return out +} + +// ----- Protocol ----- +// +// One command per line. Whitespace-separated. Values containing spaces or +// special characters are Go-quoted ("hello world", "line\nbreak"). +// +// SET [ttl_seconds] +// GET +// DEL +// KEYS +// PING +// QUIT +// +// Replies start with a single sigil and end with \n: +// + OK / simple string +// - error +// $ value (Go-quoted) +// _ nil (key not found) +// # count +// +// KEYS replies with a count line followed by N value lines. + +func handleConn(conn net.Conn, store *Store) { + defer conn.Close() + r := bufio.NewReader(conn) + w := bufio.NewWriter(conn) + + reply := func(s string) { + w.WriteString(s) + w.WriteByte('\n') + w.Flush() + } + + for { + line, err := r.ReadString('\n') + if err != nil { + return // client closed, or read error + } + line = strings.TrimRight(line, "\r\n") + if line == "" { + continue + } + + args, err := tokenize(line) + if err != nil { + reply("-" + err.Error()) + continue + } + + cmd := strings.ToUpper(args[0]) + switch cmd { + case "PING": + reply("+PONG") + + case "GET": + if len(args) != 2 { + reply("-usage: GET ") + break + } + if v, ok := store.Get(args[1]); ok { + reply("$" + strconv.Quote(v)) + } else { + reply("_") + } + + case "SET": + if len(args) < 3 || len(args) > 4 { + reply("-usage: SET [ttl_seconds]") + break + } + var ttl time.Duration + if len(args) == 4 { + n, err := strconv.Atoi(args[3]) + if err != nil || n <= 0 { + reply("-ttl must be a positive integer") + break + } + ttl = time.Duration(n) * time.Second + } + store.Set(args[1], args[2], ttl) + reply("+OK") + + case "DEL": + if len(args) != 2 { + reply("-usage: DEL ") + break + } + if store.Del(args[1]) { + reply("#1") + } else { + reply("#0") + } + + case "KEYS": + keys := store.Keys() + w.WriteString("#" + strconv.Itoa(len(keys)) + "\n") + for _, k := range keys { + w.WriteString("$" + strconv.Quote(k) + "\n") + } + w.Flush() + + case "QUIT": + reply("+BYE") + return + + default: + reply("-unknown command: " + args[0]) + } + } +} + +// tokenize splits a command line into args. A token is either a bare word +// (no whitespace) or a Go-style quoted string. This lets values contain +// spaces, newlines, etc., without inventing a new escape format. +func tokenize(line string) ([]string, error) { + var args []string + i := 0 + for i < len(line) { + // Skip whitespace. + for i < len(line) && (line[i] == ' ' || line[i] == '\t') { + i++ + } + if i >= len(line) { + break + } + + if line[i] == '"' { + // Find the matching close quote, honoring backslash escapes. + j := i + 1 + for j < len(line) { + if line[j] == '\\' && j+1 < len(line) { + j += 2 + continue + } + if line[j] == '"' { + break + } + j++ + } + if j >= len(line) { + return nil, fmt.Errorf("unterminated quoted string") + } + s, err := strconv.Unquote(line[i : j+1]) + if err != nil { + return nil, fmt.Errorf("bad quoted string: %v", err) + } + args = append(args, s) + i = j + 1 + } else { + // Bare word — read until whitespace. + j := i + for j < len(line) && line[j] != ' ' && line[j] != '\t' { + j++ + } + args = append(args, line[i:j]) + i = j + } + } + if len(args) == 0 { + return nil, fmt.Errorf("empty command") + } + return args, nil +} + +func main() { + sock := flag.String("socket", "/tmp/cubby.sock", "unix socket path") + flag.Parse() + + // Clean up any stale socket file from a previous run. + if _, err := os.Stat(*sock); err == nil { + _ = os.Remove(*sock) + } + + ln, err := net.Listen("unix", *sock) + if err != nil { + log.Fatalf("listen: %v", err) + } + _ = os.Chmod(*sock, 0600) // current user only + + store := NewStore() + + // Remove the socket file on Ctrl-C / SIGTERM. + sigs := make(chan os.Signal, 1) + signal.Notify(sigs, syscall.SIGINT, syscall.SIGTERM) + go func() { + <-sigs + log.Println("shutting down") + _ = ln.Close() + _ = os.Remove(*sock) + os.Exit(0) + }() + + log.Printf("listening on %s", *sock) + for { + conn, err := ln.Accept() + if err != nil { + return // listener closed by signal handler + } + go handleConn(conn, store) + } +}