stack.go

/*
Copyright 2020 Google LLC

Use of this source code is governed by a BSD-style
license that can be found in the LICENSE file or at
https://developers.google.com/open-source/licenses/bsd
*/

package reftable

import (
	"bytes"
	"errors"
	"fmt"
	"log"
	"math"
	"math/rand"
	"os"
	"path/filepath"
	"reflect"
	"strings"
	"time"
)

// CompactionStats holds some statistics of compaction over the
// lifetime of the stack.
type CompactionStats struct {
	Bytes uint64

	// All entries written, including from failed compaction attempts.
	EntriesWritten uint64
	Attempts       int
	Failures       int
}

// Stack is an auto-compacting stack of reftables.
type Stack struct {
	storage Storage
	cfg     Config

	// mutable
	stack              []*Reader
	merged             *Merged
	disableAutoCompact bool

	Stats CompactionStats
}

const listFileName = "tables.list"

// NewStack returns a new stack.
func NewStack(storage Storage, cfg Config) (*Stack, error) {
	if cfg.HashID == NullHashID {
		cfg.HashID = SHA1ID
	}
	switch cfg.HashID {
	case SHA1ID, SHA256ID:
	default:
		return nil, fmt.Errorf("reftable: unknown hash ID %q", cfg.HashID)
	}

	st := &Stack{
		storage: storage,
		cfg:     cfg,
	}

	if err := st.reload(true); err != nil {
		return nil, err
	}

	return st, nil
}

func (st *Stack) String() string {
	var nms []string
	for _, r := range st.stack {
		nms = append(nms, r.Name())
	}
	return fmt.Sprintf("%v", nms)
}

func (st *Stack) readNames() ([]string, error) {
	bs, err := st.storage.OpenBlockSource(listFileName)
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	defer bs.Close()

	data, err := bs.ReadBlock(0, int(bs.Size()))
	if err != nil {
		log.Printf("err %v %s %d", err, data, bs.Size())
		return nil, err
	}
	if errors.Is(err, os.ErrNotExist) {
		return nil, nil
	}
	if err != nil {
		return nil, err
	}
	lines := bytes.Split(data, []byte("\n"))

	var res []string
	for _, l := range lines {
		if len(l) > 0 {
			res = append(res, string(l))
		}
	}

	return res, nil
}

// Returns the merged stack. The stack is only valid until the next
// write, as writes may trigger reloads
func (st *Stack) Merged() *Merged {
	return st.merged
}

// Close releases file descriptors associated with this stack.
func (st *Stack) Close() {
	// Read the file list again, for closing files that were opened on windows.
	names, err := st.readNames()
	if err != nil {
		// On error, we won't remove anything.
		names = nil
	}
	nameSet := make(map[string]struct{}, len(names))
	for _, nm := range names {
		nameSet[nm] = struct{}{}
	}
	for _, r := range st.stack {
		r.Close()
		if _, ok := nameSet[r.Name()]; len(nameSet) > 0 && !ok {
			st.storage.Remove(r.Name())
		}
	}
	st.stack = nil
}

func (st *Stack) reloadOnce(names []string, reuseOpen bool) error {
	cur := map[string]*Reader{}

	for _, r := range st.stack {
		cur[r.Name()] = r
	}

	var newTables []*Reader
	defer func() {
		for _, t := range newTables {
			t.Close()
		}
	}()

	for _, name := range names {
		rd := cur[name]
		if reuseOpen && rd != nil {
			delete(cur, name)
		} else {
			bs, err := st.storage.OpenBlockSource(name)
			if err != nil {
				return err
			}

			rd, err = NewReader(bs, name)
			if err != nil {
				return fmt.Errorf("NewReader(%s): %v", name, err)
			}
		}
		newTables = append(newTables, rd)
	}

	// success. Swap.
	st.stack = newTables
	newTables = nil
	for _, old := range cur {
		old.Close()

		// On windows, we may only be able to close after
		// closing file handles.
		st.storage.Remove(old.Name())
	}
	return nil
}

// TODO: reload does not cache the (st_dev, st_ino) of tables.list,
// so it cannot distinguish "tables.list unchanged" from "tables.list
// was replaced by a file with the same inode after the inode was
// recycled". The C version (stack.c) caches list_st across reloads
// to defeat this ABA race. UpToDate has the same exposure.
func (st *Stack) reload(reuseOpen bool) error {
	var delay time.Duration
	deadline := time.Now().Add(5 * time.Second / 2)
	for time.Now().Before(deadline) {
		names, err := st.readNames()
		if err != nil {
			return err
		}
		err = st.reloadOnce(names, reuseOpen)
		if err == nil {
			break
		}
		if !errors.Is(err, os.ErrNotExist) {
			return err
		}
		after, err := st.readNames()
		if err != nil {
			return err
		}
		if reflect.DeepEqual(after, names) {
			// XXX: propogate name
			return os.ErrNotExist
		}

		// compaction changed names; back off and retry.
		delay = 2*delay + time.Millisecond*time.Duration(1+rand.Intn(2))
		time.Sleep(delay)
	}

	var tabs []Table
	for _, r := range st.stack {
		tabs = append(tabs, r)
	}

	m, err := NewMerged(tabs, st.cfg.HashID)
	if err != nil {
		return err
	}
	m.suppressDeletions = true
	st.merged = m
	return nil
}

// ErrLockFailure is returned for failed writes. On a failed write,
// the stack is reloaded, so the transaction may be retried.
var ErrLockFailure = errors.New("reftable: lock failure")

func (st *Stack) UpToDate() (bool, error) {
	names, err := st.readNames()
	if err != nil {
		return false, err
	}

	if len(names) != len(st.stack) {
		return false, nil
	}

	for i, e := range st.stack {
		if e.name != names[i] {
			return false, nil
		}
	}
	return true, nil
}

// Add a new reftable to stack, transactionally.
func (st *Stack) Add(write func(w *Writer) error) error {
	if err := st.add(write); err != nil {
		if err == ErrLockFailure {
			st.reload(true)
		}
		return err
	}

	if !st.disableAutoCompact {
		return st.AutoCompact()
	}
	return nil
}

func (st *Stack) add(write func(w *Writer) error) error {
	tr, err := st.NewAddition()
	if err != nil {
		return err
	}
	defer tr.Close()
	if err := tr.Add(write); err != nil {
		return err
	}

	return tr.Commit()
}

// Addition is a transaction that adds new tables to the top of the
// stack.
type Addition struct {
	lockFile        AtomicWriter
	stack           *Stack
	names           []string
	newTables       []string
	nextUpdateIndex uint64
}

// NewAddition returns an Addition instance. As a side effect, this
// takes a global filesystem lock on the ref database.
func (st *Stack) NewAddition() (*Addition, error) {
	tr := Addition{
		stack: st,
	}
	var err error
	tr.lockFile, err = st.storage.LockForWrite(listFileName)
	if errors.Is(err, os.ErrExist) {
		return nil, ErrLockFailure
	}
	if err != nil {
		return nil, err
	}
	for _, e := range st.stack {
		tr.names = append(tr.names, e.name)
	}
	if ok, err := tr.stack.UpToDate(); err != nil {
		tr.Close()
		return nil, err
	} else if !ok {
		tr.Close()
		return nil, ErrLockFailure
	}
	tr.nextUpdateIndex = tr.stack.NextUpdateIndex()
	return &tr, nil
}

// Add calls the given function to write a new table at the top of
// the stack.
func (tr *Addition) Add(write func(w *Writer) error) error {
	dest := formatName(tr.nextUpdateIndex, tr.nextUpdateIndex) + ".ref"

	tab, err := tr.stack.storage.Update(dest)
	if err != nil {
		return err
	}
	defer func() {
		if tab != nil {
			tab.Close()
		}
	}()

	wr, err := NewWriter(tab, &tr.stack.cfg)
	if err != nil {
		return err
	}

	if err := write(wr); err != nil {
		return err
	}

	if err := wr.Close(); err != nil {
		if errors.Is(err, ErrEmptyTable) {
			return nil
		}
		return err
	}

	if wr.minUpdateIndex < tr.nextUpdateIndex {
		return ErrLockFailure
	}

	if err := tr.stack.checkAddition(filepath.Base(tab.Name())); err != nil {
		return err
	}

	if err := tab.Commit(); err != nil {
		return err
	}

	tr.names = append(tr.names, dest)
	tr.newTables = append(tr.newTables, dest)
	tr.nextUpdateIndex = wr.maxUpdateIndex + 1
	return nil
}

// Close releases all non-committed data from the transaction.
func (tr *Addition) Close() {
	for _, nm := range tr.newTables {
		tr.stack.storage.Remove(nm)
	}
	tr.lockFile.Close()
}

// Commit commits the changes to the database, releasing the lock.
func (tr *Addition) Commit() error {
	if len(tr.newTables) == 0 {
		// Nothing to be done.
		return nil
	}

	if _, err := tr.lockFile.Write([]byte(strings.Join(tr.names, "\n"))); err != nil {
		tr.Close()
		return err
	}

	if err := tr.lockFile.Commit(); err != nil {
		tr.Close()
		return err
	}
	tr.newTables = nil

	return tr.stack.reload(true)
}

func (s *Stack) checkAddition(tabname string) error {
	if s.cfg.SkipNameCheck {
		return nil
	}
	bs, err := s.storage.OpenBlockSource(tabname)
	if err != nil {
		return err
	}
	r, err := NewReader(bs, tabname)
	if err != nil {
		return err
	}
	defer r.Close()
	it, err := r.SeekRef("")
	if err != nil {
		return err
	}

	var recs []RefRecord
	for {
		var rec RefRecord
		ok, err := it.NextRef(&rec)
		if err != nil {
			return err
		}
		if !ok {
			break
		}
		recs = append(recs, rec)
	}

	return validateRefRecordAddition(s.Merged(), recs)
}

// non-deterministic random generator.
var randomRandom = rand.New(rand.NewSource(time.Now().UnixNano()))

func formatName(min, max uint64) string {
	return fmt.Sprintf("0x%012x-0x%012x-%08x", min, max, randomRandom.Uint32())
}

// NextUpdateIndex returns the update index at which to write the next table.
func (st *Stack) NextUpdateIndex() uint64 {
	if sz := len(st.stack); sz > 0 {
		return st.stack[sz-1].MaxUpdateIndex() + 1
	}
	return 1
}

// compactLocked writes the compacted version of tables [first,last]
// into a temporary file, whose name is returned.
func (st *Stack) compactLocked(first, last int, expiration *LogExpirationConfig) (AtomicWriter, error) {
	fn := formatName(st.stack[first].MinUpdateIndex(),
		st.stack[last].MaxUpdateIndex()) + ".ref"

	tmpTable, err := st.storage.Update(fn)
	if err != nil {
		return nil, err
	}
	defer func() {
		if tmpTable != nil {
			tmpTable.Close()
		}
	}()

	wr, err := NewWriter(tmpTable, &st.cfg)
	if err != nil {
		return nil, err
	}

	if err := st.writeCompact(wr, first, last, expiration); err != nil {
		return nil, err
	}

	if err := wr.Close(); err != nil {
		return nil, err
	}

	result := tmpTable
	tmpTable = nil
	return result, nil
}

func (st *Stack) writeCompact(wr *Writer, first, last int, expiration *LogExpirationConfig) error {
	// do it.
	wr.SetLimits(st.stack[first].MinUpdateIndex(),
		st.stack[last].MaxUpdateIndex())

	var subtabs []Table
	for i := first; i <= last; i++ {
		subtabs = append(subtabs, st.stack[i])
	}

	merged, err := NewMerged(subtabs, st.cfg.HashID)
	if err != nil {
		return err
	}
	it, err := merged.SeekRef("")
	if err != nil {
		return err
	}

	var entries uint64
	for {
		var rec RefRecord
		ok, err := it.NextRef(&rec)
		if err != nil {
			return err
		}
		if !ok {
			break
		}

		if first == 0 && rec.IsDeletion() {
			continue
		}

		if err := wr.AddRef(&rec); err != nil {
			return err
		}
		entries++
	}

	it, err = merged.SeekLog("", math.MaxUint64)
	if err != nil {
		return err
	}
	for {
		var rec LogRecord
		ok, err := it.NextLog(&rec)
		if err != nil {
			return err
		}
		if !ok {
			break
		}

		if expiration != nil {
			if expiration.Time > 0 && rec.Time < expiration.Time {
				continue
			}

			if expiration.MaxUpdateIndex != 0 && rec.UpdateIndex > expiration.MaxUpdateIndex {
				continue
			}
			if expiration.MinUpdateIndex != 0 && rec.UpdateIndex < expiration.MinUpdateIndex {
				continue
			}
		}

		if err := wr.AddLog(&rec); err != nil {
			return err
		}
		entries++
	}

	st.Stats.EntriesWritten += entries
	return nil
}

func (st *Stack) compactRangeStats(first, last int, expiration *LogExpirationConfig) (bool, error) {
	ok, err := st.compactRange(first, last, expiration)
	if !ok {
		st.Stats.Failures++
	}
	return ok, err
}

func (st *Stack) compactRange(first, last int, expiration *LogExpirationConfig) (bool, error) {
	if first >= last && expiration == nil {
		return true, nil
	}
	st.Stats.Attempts++

	lock, err := st.storage.LockForWrite(listFileName)
	if errors.Is(err, os.ErrExist) {
		return false, nil
	}
	if err != nil {
		return false, err
	}

	defer func() {
		lock.Close()
	}()

	if ok, err := st.UpToDate(); !ok || err != nil {
		return false, err
	}

	var deleteOnSuccess []string
	var subtableLocks []AtomicWriter
	defer func() {
		for _, l := range subtableLocks {
			l.Close()
		}
	}()
	for i := first; i <= last; i++ {
		subtab := st.stack[i].name
		subtabLock, err := st.storage.LockForWrite(subtab)

		if errors.Is(err, os.ErrExist) {
			return false, nil
		}
		if err != nil {
			return false, err
		}
		subtableLocks = append(subtableLocks, subtabLock)
		deleteOnSuccess = append(deleteOnSuccess, subtab)
	}

	lock.Close()

	tmpTable, err := st.compactLocked(first, last, expiration)

	// Compaction + tombstones can create an empty table out of non-empty tables.
	if errors.Is(err, ErrEmptyTable) {
		// In this case, we may have tmpTable == nil
		err = nil
	}

	if err != nil {
		return false, err
	}
	if tmpTable != nil {
		defer tmpTable.Close()
	}

	lock, err = st.storage.LockForWrite(listFileName)
	if err != nil {
		return false, err
	}

	defer lock.Close()

	if tmpTable != nil {
		if err := tmpTable.Commit(); err != nil {
			return false, err
		}
	}

	var names []string
	for i := range first {
		names = append(names, st.stack[i].name)
	}

	if tmpTable != nil {
		names = append(names, tmpTable.Name())
	}

	for i := last + 1; i < len(st.stack); i++ {
		names = append(names, st.stack[i].name)
	}

	if _, err := lock.Write([]byte(strings.Join(names, "\n"))); err != nil {
		if tmpTable != nil {
			os.Remove(tmpTable.Name())
		}
		return false, err
	}
	if err := lock.Commit(); err != nil {
		if tmpTable != nil {
			os.Remove(tmpTable.Name())
		}
		return false, err
	}

	for _, nm := range deleteOnSuccess {
		if tmpTable != nil && nm != tmpTable.Name() {
			// reflog expiry might cause us to reopen a
			// new file with the same name.
			os.Remove(nm)
		}
	}

	// If we expire log entries on a full compaction we write a
	// table with the same the (min,max) update index, but we have
	// to read from disk again.
	if err := st.reload(expiration == nil); err != nil {
		return true, fmt.Errorf("reload: %w", err)
	}
	return true, err
}

func (st *Stack) tableSizesForCompaction() []uint64 {
	var res []uint64
	version := 1
	if st.cfg.HashID == SHA256ID {
		version = 2
	}
	var overhead = uint64(headerSize(version) - 1)
	for _, t := range st.stack {
		res = append(res, t.size-overhead)
	}
	return res
}

type segment struct {
	start int
	end   int // exclusive
	log   int
	bytes uint64
}

func (st *segment) size() int { return st.end - st.start }

func log2(sz uint64) int {
	base := uint64(2)
	if sz == 0 {
		return 0
	}

	l := 0
	for sz > 0 {
		l++
		sz /= base
	}

	return l - 1
}

func sizesToSegments(sizes []uint64) []segment {
	var cur segment
	var res []segment
	for i, sz := range sizes {
		l := log2(sz)
		if cur.log != l && cur.bytes > 0 {
			res = append(res, cur)
			cur = segment{
				start: i,
			}
		}
		cur.log = l
		cur.end = i + 1
		cur.bytes += sz
	}

	res = append(res, cur)
	return res
}

/*
We play the game of 2048: consecutive tables of the same size (as
determined by their log2) are compacted together. We try to combine
the result with preceding tables, if they are smaller (as determined
by their log2). As a result, if we have N entries, each entry will
go into a bigger table in a maximum of log2(N) times, making for
log2(N) * N overall cost.
*/
func suggestCompactionSegment(sizes []uint64) *segment {
	segs := sizesToSegments(sizes)

	minSeg := segment{log: 64}
	for _, st := range segs {
		if st.size() == 1 {
			continue
		}

		if st.log < minSeg.log {
			minSeg = st
		}
	}
	if minSeg.size() == 0 {
		return nil
	}

	for minSeg.start > 0 {
		prev := minSeg.start - 1
		if log2(minSeg.bytes) < log2(sizes[prev]) {
			break
		}

		minSeg.start = prev
		minSeg.bytes += sizes[prev]
	}

	return &minSeg
}

// AutoCompact runs a compaction if the stack looks imbalanced.
func (st *Stack) AutoCompact() error {
	sizes := st.tableSizesForCompaction()
	seg := suggestCompactionSegment(sizes)
	if seg != nil {
		_, err := st.compactRangeStats(seg.start, seg.end-1, nil)
		return err
	}
	return nil
}

// CompactAll compacts the entire stack. If expiration is given, expire log entries.
func (st *Stack) CompactAll(expiration *LogExpirationConfig) error {
	_, err := st.compactRange(0, len(st.stack)-1, expiration)
	return err
}

// Clean removes stale *.ref files. It is only required to be called
// on Windows, if previous processes did not call Stack.Close on exit.
func (st *Stack) Clean() error {
	// Take a lock to prevent concurrent updates.
	add, err := st.NewAddition()
	if err != nil {
		return err
	}
	defer add.Close()

	if err := st.reload(true); err != nil {
		return err
	}

	names := map[string]struct{}{}
	for _, r := range st.stack {
		names[r.Name()] = struct{}{}
	}
	entries, err := st.storage.ReadDir()
	if err != nil {
		return err
	}

	max := st.merged.MaxUpdateIndex()
	for _, e := range entries {
		name := e.Name()
		if _, ok := names[name]; ok {
			continue
		}
		if !strings.HasSuffix(name, ".ref") {
			continue
		}

		bs, err := st.storage.OpenBlockSource(name)
		if err != nil {
			return err
		}

		rd, err := NewReader(bs, name)
		if err != nil {
			return fmt.Errorf("NewReader(%s): %v", name, err)
		}

		cur := rd.MaxUpdateIndex()
		rd.Close()
		if cur <= max {
			st.storage.Remove(name)
		}
	}
	return nil
}