package storage import ( "bufio" "fmt" "io" "os" "path/filepath" "strings" "sync" "syscall" ) type ValueError struct { Value any } func (err ValueError) Error() string { return fmt.Sprintf("kv: can't store value of type %T", err.Value) } type readSeekerKV struct { mu sync.Mutex r io.ReadSeeker c io.Closer sep rune } func (kv *readSeekerKV) scan(key string) (data string, ok bool, err error) { kv.mu.Lock() if _, err = kv.r.Seek(0, io.SeekStart); err != nil { return } s := bufio.NewScanner(kv.r) for s.Scan() { line := s.Text() if i := strings.IndexRune(line, kv.sep); i > -1 && line[:i] == key { kv.mu.Unlock() return line[i+1:], true, nil } } err = s.Err() kv.mu.Unlock() return } func (kv *readSeekerKV) Has(key string) bool { _, ok, err := kv.scan(key) return ok && err == nil } func (kv *readSeekerKV) Get(key string) (value any, ok bool) { value, ok, _ = kv.scan(key) return } func (kv *readSeekerKV) Close() error { if kv.c == nil { return nil } return kv.c.Close() } func OpenKV(name string, sep rune) (KV, error) { f, err := os.Open(name) if err != nil { return nil, err } return &readSeekerKV{ r: f, sep: sep, }, nil } type readSeekWriterKV struct { readSeekerKV w io.WriteSeeker } type truncater interface { Truncate(int64) error } func (kv *readSeekWriterKV) Set(key string, value any) error { var line string switch value := value.(type) { case string: line = key + string(kv.sep) + value + "\n" case []string: line = key + string(kv.sep) + strings.Join(value, string(kv.sep)) + "\n" default: return ValueError{Value: value} } kv.mu.Lock() defer kv.mu.Unlock() if _, err := kv.r.Seek(0, io.SeekStart); err != nil { return err } var ( lines []string found bool scanner = bufio.NewScanner(kv.r) ) for scanner.Scan() { text := scanner.Text() if i := strings.IndexRune(text, kv.sep); i > -1 { if found = text[:i] == key; found { lines = append(lines, line) } else { lines = append(lines, text) } } else { lines = append(lines, line) } } if !found { lines = append(lines, line) } // Writing strategy: if it's a file, write to a new file and move it over. if f, ok := kv.w.(*os.File); ok { return kv.replaceFile(f, lines) } if t, ok := kv.w.(truncater); ok { if err := t.Truncate(0); err != nil { return err } } if _, err := kv.w.Seek(0, io.SeekStart); err != nil { return err } for _, line := range lines { if _, err := io.WriteString(kv.w, line+"\n"); err != nil { return err } } return nil } func (kv *readSeekWriterKV) replaceFile(f *os.File, lines []string) (err error) { var ( prev = f.Name() info os.FileInfo ) if info, err = f.Stat(); err != nil { return } var ( name = "." + filepath.Base(f.Name()) + ".*" n *os.File ) if n, err = os.CreateTemp(filepath.Dir(f.Name()), name); err != nil { return } name = n.Name() // Replicate original file mode if err = os.Chmod(name, info.Mode()); err != nil { _ = n.Close() _ = os.Remove(name) return } // Replicate original file ownership if stat, ok := info.Sys().(*syscall.Stat_t); ok { // This may fail if we aren't allowed, ignore it. _ = os.Chown(name, int(stat.Uid), int(stat.Gid)) } // Write lines to tempfile. for _, line := range lines { if _, err = io.WriteString(n, line+"\n"); err != nil { _ = n.Close() _ = os.Remove(name) return } } if err = n.Close(); err != nil { _ = os.Remove(name) return } // Close original file and replace it. _ = f.Close() if err = os.Rename(name, prev); err != nil { return } // Reopen file and replace our readers/writers/closers if f, err = os.OpenFile(prev, os.O_APPEND|os.O_RDWR, info.Mode()|os.ModePerm); err != nil { return } kv.r = f kv.w = f kv.c = f return } func OpenWritableKV(name string, sep rune, perm os.FileMode) (WritableKV, error) { f, err := os.OpenFile(name, os.O_CREATE|os.O_RDWR, perm) if err != nil { return nil, err } return &readSeekWriterKV{ readSeekerKV: readSeekerKV{ r: f, c: f, sep: sep, }, w: f, }, nil }