conduit/storage/io.go

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
}