Initial import

storage/io.go

@@ -0,0 +1,225 @@
+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
+}