litestream/replicator.go

package litestream

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

// Replicator represents a method for replicating the snapshot & WAL data to
// a remote destination.
type Replicator interface {
	Name() string
	Type() string
	Snapshotting() bool
	Start(ctx context.Context)
	Stop()
}

var _ Replicator = (*FileReplicator)(nil)

// FileReplicator is a replicator that replicates a DB to a local file path.
type FileReplicator struct {
	db   *DB    // source database
	name string // replicator name, optional
	dst  string // destination path

	mu           sync.RWMutex
	wg           sync.WaitGroup
	snapshotting bool // if true, currently copying database

	ctx    context.Context
	cancel func()
}

// NewFileReplicator returns a new instance of FileReplicator.
func NewFileReplicator(db *DB, name, dst string) *FileReplicator {
	return &FileReplicator{
		db:     db,
		name:   name,
		dst:    dst,
		cancel: func() {},
	}
}

// Name returns the name of the replicator. Returns the type if no name set.
func (r *FileReplicator) Name() string {
	if r.name != "" {
		return r.name
	}
	return r.Type()
}

// Type returns the type of replicator.
func (r *FileReplicator) Type() string {
	return "file"
}

// SnapshotPath returns the path to a snapshot file.
func (r *FileReplicator) SnapshotPath(generation string, index int) string {
	return filepath.Join(r.dst, "generations", generation, "snapshots", fmt.Sprintf("%016x.snapshot", index))
}

// WALPath returns the path to a WAL file.
func (r *FileReplicator) WALPath(generation string, index int) string {
	return filepath.Join(r.dst, "generations", generation, "wal", fmt.Sprintf("%016x.wal", index))
}

// Snapshotting returns true if replicator is current snapshotting.
func (r *FileReplicator) Snapshotting() bool {
	r.mu.RLock()
	defer r.mu.RLock()
	return r.snapshotting
}

// Start starts replication for a given generation.
func (r *FileReplicator) Start(ctx context.Context) {
	// Stop previous replication.
	r.Stop()

	r.mu.Lock()
	defer r.mu.Unlock()

	// Set snapshotting state.
	r.snapshotting = true

	// Wrap context with cancelation.
	ctx, r.cancel = context.WithCancel(ctx)

	// Start goroutine to replicate data.
	r.wg.Add(1)
	go func() { defer r.wg.Done(); r.monitor(ctx) }()
}

// Stop cancels any outstanding replication and blocks until finished.
func (r *FileReplicator) Stop() {
	r.cancel()
	r.wg.Wait()
}

// monitor runs in a separate goroutine and continuously replicates the DB.
func (r *FileReplicator) monitor(ctx context.Context) {
	// Continuously check for new data to replicate.
	ch := make(chan struct{})
	close(ch)
	var notify <-chan struct{} = ch

	var pos Pos
	var err error
	for {
		select {
		case <-ctx.Done():
			return
		case <-notify:
		}

		// Fetch new notify channel before replicating data.
		notify = r.db.Notify()

		// Determine position, if necessary.
		if pos.IsZero() {
			if pos, err = r.pos(); err != nil {
				log.Printf("%s(%s): cannot determine position: %w", r.db.Path(), r.Name(), err)
				continue
			} else if pos.IsZero() {
				log.Printf("%s(%s): no generation, waiting for data", r.db.Path(), r.Name())
				continue
			}
		}

		// If we have no replicated WALs, start from last index in shadow WAL.
		if pos.Index == 0 && pos.Offset == 0 {
			if pos.Index, err = r.db.CurrentShadowWALIndex(pos.Generation); err != nil {
				log.Printf("%s(%s): cannot determine latest shadow wal index: %s", r.db.Path(), r.Name(), err)
				continue
			}
		}

		// Synchronize the shadow wal into the replication directory.
		if pos, err = r.sync(ctx, pos); err != nil {
			log.Printf("%s(%s): sync error: %s", r.db.Path(), r.Name(), err)
			continue
		}
	}
}

// pos returns the position for the replicator for the current generation.
// Returns a zero value if there is no active generation.
func (r *FileReplicator) pos() (pos Pos, err error) {
	// Find the current generation from the DB. Return zero pos if no generation.
	generation, err := r.db.CurrentGeneration()
	if err != nil {
		return pos, err
	} else if generation == "" {
		return pos, nil // empty position
	}
	pos.Generation = generation

	// Find the max WAL file.
	walDir := filepath.Join(r.dst, "generations", generation, "wal")
	fis, err := ioutil.ReadDir(walDir)
	if os.IsNotExist(err) {
		return pos, nil // no replicated wal, start at beginning of generation
	} else if err != nil {
		return pos, err
	}

	index := -1
	for _, fi := range fis {
		if !strings.HasSuffix(fi.Name(), WALExt) {
			continue
		}

		if v, err := ParseWALFilename(filepath.Base(fi.Name())); err != nil {
			continue // invalid wal filename
		} else if index == -1 || v > index {
			index = v
		}
	}
	if index == -1 {
		return pos, nil // wal directory exists but no wal files, return beginning pos
	}
	pos.Index = index

	// Determine current offset.
	fi, err := os.Stat(filepath.Join(walDir, FormatWALFilename(pos.Index)))
	if err != nil {
		return pos, err
	}
	pos.Offset = fi.Size()

	return pos, nil
}

// snapshot copies the entire database to the replica path.
func (r *FileReplicator) snapshot(ctx context.Context, generation string, index int) error {
	// Mark replicator as snapshotting to prevent checkpoints by the DB.
	r.mu.Lock()
	r.snapshotting = true
	r.mu.Unlock()

	// Ensure we release the snapshot flag when we leave the function.
	defer func() {
		r.mu.Lock()
		r.snapshotting = false
		r.mu.Unlock()
	}()

	// Ignore if we already have a snapshot for the given WAL index.
	snapshotPath := r.SnapshotPath(generation, index)
	if _, err := os.Stat(snapshotPath); err == nil {
		return nil
	}

	rd, err := os.Open(r.db.Path())
	if err != nil {
		return err
	}
	defer rd.Close()

	if err := os.MkdirAll(filepath.Dir(snapshotPath), 0700); err != nil {
		return err
	}

	w, err := os.Create(snapshotPath + ".tmp")
	if err != nil {
		return err
	}
	defer w.Close()

	if _, err := io.Copy(w, rd); err != nil {
		return err
	} else if err := w.Sync(); err != nil {
		return err
	} else if err := w.Close(); err != nil {
		return err
	} else if err := os.Rename(snapshotPath+".tmp", snapshotPath); err != nil {
		return err
	}

	return nil
}

// snapshotN returns the number of snapshots for a generation.
func (r *FileReplicator) snapshotN(generation string) (int, error) {
	fis, err := ioutil.ReadDir(filepath.Join(r.dst, "generations", generation, "snapshots"))
	if os.IsNotExist(err) {
		return 0, nil
	} else if err != nil {
		return 0, err
	}

	var n int
	for _, fi := range fis {
		if strings.HasSuffix(fi.Name(), SnapshotExt) {
			n++
		}
	}
	return n, nil
}

func (r *FileReplicator) sync(ctx context.Context, pos Pos) (_ Pos, err error) {
	// Read all WAL files since the last position.
	for {
		if pos, err = r.syncNext(ctx, pos); err == io.EOF {
			return pos, nil
		} else if err != nil {
			return pos, err
		}
	}
}

func (r *FileReplicator) syncNext(ctx context.Context, pos Pos) (_ Pos, err error) {
	rd, err := r.db.WALReader(pos)
	if err == io.EOF {
		return pos, err
	} else if err != nil {
		return pos, fmt.Errorf("wal reader: %w", err)
	}
	defer rd.Close()

	// Create snapshot if no snapshots exist.
	if n, err := r.snapshotN(rd.Pos().Generation); err != nil {
		return pos, err
	} else if n == 0 {
		if err := r.snapshot(ctx, rd.Pos().Generation, rd.Pos().Index); err != nil {
			return pos, err
		}
	}

	// Ensure parent directory exists for WAL file.
	filename := r.WALPath(rd.Pos().Generation, rd.Pos().Index)
	if err := os.MkdirAll(filepath.Dir(filename), 0700); err != nil {
		return pos, err
	}

	// Create a temporary file to write into so we don't have partial writes.
	w, err := os.OpenFile(filename, os.O_RDWR|os.O_CREATE, 0600)
	if err != nil {
		return pos, err
	}
	defer w.Close()

	// Seek, copy & sync WAL contents.
	if _, err := w.Seek(rd.Pos().Offset, io.SeekStart); err != nil {
		return pos, err
	} else if _, err := io.Copy(w, rd); err != nil {
		return pos, err
	} else if err := w.Sync(); err != nil {
		return pos, err
	} else if err := w.Close(); err != nil {
		return pos, err
	}

	// Return ending position of the reader.
	return rd.Pos(), nil
}