1package snapshot
  2
  3import (
  4	"errors"
  5	"fmt"
  6	"os"
  7	"path/filepath"
  8)
  9
 10type atomicFileWriter struct {
 11	path string
 12	tmp  string
 13	f    *os.File
 14}
 15
 16var (
 17	ErrEmptyPath    = errors.New("empty path")
 18	ErrWriterClosed = errors.New("write after close")
 19)
 20
 21const (
 22	_dirMode = 0o750
 23)
 24
 25func newAtomicFileWriter(path string) (*atomicFileWriter, error) {
 26	if path == "" {
 27		return nil, ErrEmptyPath
 28	}
 29
 30	dir := filepath.Dir(path)
 31	err := os.MkdirAll(dir, _dirMode)
 32	if err != nil {
 33		err = fmt.Errorf("mkdir %s: %w", dir, err)
 34		return nil, err
 35	}
 36
 37	tmp, err := os.CreateTemp(dir, ".snapshot-tmp-*")
 38	if err != nil {
 39		err = fmt.Errorf("create temp file in %s: %w", dir, err)
 40		return nil, err
 41	}
 42
 43	return &atomicFileWriter{
 44		path: path,
 45		tmp:  tmp.Name(),
 46		f:    tmp,
 47	}, nil
 48}
 49
 50func (w *atomicFileWriter) Write(p []byte) (int, error) {
 51	if w.f == nil {
 52		return 0, ErrWriterClosed
 53	}
 54	return w.f.Write(p)
 55}
 56
 57type aborter interface{ Abort() }
 58
 59// Abort is an alternate Close which omits flushing and renaming to w.path
 60func (w *atomicFileWriter) Abort() {
 61	if w.f == nil {
 62		return
 63	}
 64
 65	f := w.f
 66	tmp := w.tmp
 67	w.f = nil
 68	_ = f.Close()
 69	_ = os.Remove(tmp)
 70}
 71
 72func (w *atomicFileWriter) Close() error {
 73	if w.f == nil {
 74		// already closed
 75		return nil
 76	}
 77
 78	// Capture state, nil out to prevent reuse
 79	f := w.f
 80	tmp := w.tmp
 81	path := w.path
 82	w.f = nil
 83
 84	err := f.Sync()
 85	if err != nil {
 86		_ = f.Close()
 87		_ = os.Remove(tmp)
 88		return fmt.Errorf("syncing %s: %w", tmp, err)
 89	}
 90	err = f.Close()
 91	if err != nil {
 92		_ = os.Remove(tmp)
 93		return fmt.Errorf("closing %s: %w", tmp, err)
 94	}
 95
 96	err = os.Rename(tmp, path)
 97	if err != nil {
 98		_ = os.Remove(tmp)
 99		return fmt.Errorf("renaming %s -> %s: %w", tmp, path, err)
100	}
101
102	return nil
103}