github.com/zuoyebang/bitalostable@v1.0.1-0.20240229032404-e3b99a834294/vfs/disk_full.go

// Copyright 2021 The LevelDB-Go and Pebble Authors. All rights reserved. Use
// of this source code is governed by a BSD-style license that can be found in
// the LICENSE file.

package vfs

import (
	"io"
	"os"
	"sync"
	"sync/atomic"
	"syscall"

	"github.com/cockroachdb/errors"
)

// OnDiskFull wraps the provided FS with an FS that examines returned errors,
// looking for ENOSPC errors. It invokes the provided callback when the
// underlying filesystem returns an error signifying the storage is out of
// disk space.
//
// All new writes to the filesystem are blocked while the callback executes,
// so care must be taken to avoid expensive work from within the callback.
//
// Once the callback completes, any write-oriented operations that encountered
// ENOSPC are retried exactly once. Once the callback completes, it will not
// be invoked again until a new operation that began after the callback
// returned encounters an ENOSPC error.
//
// OnDiskFull may be used to automatically manage a ballast file, which is
// removed from the filesystem from within the callback. Note that if managing
// a ballast, the caller should maintain a reference to the inner FS and
// remove the ballast on the unwrapped FS.
func OnDiskFull(fs FS, fn func()) FS {
	newFS := &enospcFS{inner: fs}
	newFS.mu.Cond.L = &newFS.mu.Mutex
	newFS.mu.onDiskFull = fn
	return newFS
}

type enospcFS struct {
	inner  FS
	atomic struct {
		// generation is a monotonically increasing number that encodes the
		// current state of ENOSPC error handling. Incoming writes are
		// organized into generations to provide strong guarantees on when the
		// disk full callback is invoked. The callback is invoked once per
		// write generation.
		//
		// Special significance is given to the parity of this generation
		// field to optimize incoming writes in the normal state, which only
		// need to perform a single atomic load. If generation is odd, an
		// ENOSPC error is being actively handled. The generations associated
		// with writes are always even.
		//
		// The lifecycle of a write is:
		//
		// 1. Atomically load the current generation.
		//    a. If it's even, this is the write's generation number.
		//    b. If it's odd, an ENOSPC was recently encountered and the
		//       corresponding invocation of the disk full callback has not
		//       yet completed. The write must wait until the callback has
		//       completed and generation is updated to an even number, which
		//       becomes the write's generation number.
		// 2. Perform the write. If it encounters no error or an error other
		//    than ENOSPC, the write returns and proceeds no further in this
		//    lifecycle.
		// 3. Handle ENOSPC. If the write encounters ENOSPC, the callback must
		//    be invoked for the write's generation. The write's goroutine
		//    acquires the FS's mutex.
		//    a. If the FS's current generation is still equal to the write's
		//       generation, the write is the first write of its generation to
		//       encounter ENOSPC. It increments the FS's current generation
		//       to an odd number, signifying that an ENOSPC is being handled
		//       and invokes the callback.
		//    b. If the FS's current generation has changed, some other write
		//       from the same generation encountered an ENOSPC first. This
		//       write waits on the condition variable until the FS's current
		//       generation is updated indicating that the generation's
		//       callback invocation has completed.
		// 3. Retry the write once. The callback for the write's generation
		//    has completed, either by this write's goroutine or another's.
		//    The write may proceed with the expectation that the callback
		//    remedied the full disk by freeing up disk space and an ENOSPC
		//    should not be encountered again for at least a few minutes. If
		//    we do encounter another ENOSPC on the retry, the callback was
		//    unable to remedy the full disk and another retry won't be
		//    useful. Any error, including ENOSPC, during the retry is
		//    returned without further handling.  None of the retries invoke
		//    the callback.
		//
		// This scheme has a few nice properties:
		// * Once the disk-full callback completes, it won't be invoked
		//   again unless a write that started strictly later encounters an
		//   ENOSPC. This is convenient if the callback strives to 'fix' the
		//   full disk, for example, by removing a ballast file. A new
		//   invocation of the callback guarantees a new problem.
		// * Incoming writes block if there's an unhandled ENOSPC. Some
		//   writes, like WAL or MANIFEST fsyncs, are fatal if they encounter
		//   an ENOSPC.
		generation uint32
	}
	mu struct {
		sync.Mutex
		sync.Cond
		onDiskFull func()
	}
}

// Unwrap returns the underlying FS. This may be called by vfs.Root to access
// the underlying filesystem.
func (fs *enospcFS) Unwrap() FS {
	return fs.inner
}

// waitUntilReady is called before every FS or File operation that
// might return ENOSPC. If an ENOSPC was encountered and the corresponding
// invocation of the `onDiskFull` callback has not yet returned,
// waitUntilReady blocks until the callback returns. The returned generation
// is always even.
func (fs *enospcFS) waitUntilReady() uint32 {
	gen := atomic.LoadUint32(&fs.atomic.generation)
	if gen%2 == 0 {
		// An even generation indicates that we're not currently handling an
		// ENOSPC. Allow the write to proceed.
		return gen
	}

	// We're currently handling an ENOSPC error. Wait on the condition
	// variable until we're not handling an ENOSPC.
	fs.mu.Lock()
	defer fs.mu.Unlock()

	// Load the generation again with fs.mu locked.
	gen = atomic.LoadUint32(&fs.atomic.generation)
	for gen%2 == 1 {
		fs.mu.Wait()
		gen = atomic.LoadUint32(&fs.atomic.generation)
	}
	return gen
}

func (fs *enospcFS) handleENOSPC(gen uint32) {
	fs.mu.Lock()
	defer fs.mu.Unlock()

	currentGeneration := atomic.LoadUint32(&fs.atomic.generation)

	// If the current generation is still `gen`, this is the first goroutine
	// to hit an ENOSPC within this write generation, so this goroutine is
	// responsible for invoking the callback.
	if currentGeneration == gen {
		// Increment the generation to an odd number, indicating that the FS
		// is out-of-disk space and incoming writes should pause and wait for
		// the next generation before continuing.
		atomic.StoreUint32(&fs.atomic.generation, gen+1)

		func() {
			// Drop the mutex while we invoke the callback, re-acquiring
			// afterwards.
			fs.mu.Unlock()
			defer fs.mu.Lock()
			fs.mu.onDiskFull()
		}()

		// Update the current generation again to an even number, indicating
		// that the callback has completed for the write generation `gen`.
		atomic.StoreUint32(&fs.atomic.generation, gen+2)
		fs.mu.Broadcast()
		return
	}

	// The current generation has already been incremented, so either the
	// callback is currently being run by another goroutine or it's already
	// completed. Wait for it complete if it hasn't already.
	//
	// The current generation may be updated multiple times, including to an
	// odd number signifying a later write generation has already encountered
	// ENOSPC. In that case, the callback was not able to remedy the full disk
	// and waiting is unlikely to be helpful.  Continuing to wait risks
	// blocking an unbounded number of generations.  Retrying and bubbling the
	// ENOSPC up might be helpful if we can abort a large compaction that
	// started before we became more selective about compaction picking, so
	// this loop only waits for this write generation's callback and no
	// subsequent generations' callbacks.
	for currentGeneration == gen+1 {
		fs.mu.Wait()
		currentGeneration = atomic.LoadUint32(&fs.atomic.generation)
	}
}

func (fs *enospcFS) Create(name string) (File, error) {
	gen := fs.waitUntilReady()

	f, err := fs.inner.Create(name)

	if err != nil && isENOSPC(err) {
		fs.handleENOSPC(gen)
		f, err = fs.inner.Create(name)
	}
	if f != nil {
		f = WithFd(f, enospcFile{
			fs:    fs,
			inner: f,
		})
	}
	return f, err
}

func (fs *enospcFS) Link(oldname, newname string) error {
	gen := fs.waitUntilReady()

	err := fs.inner.Link(oldname, newname)

	if err != nil && isENOSPC(err) {
		fs.handleENOSPC(gen)
		err = fs.inner.Link(oldname, newname)
	}
	return err
}

func (fs *enospcFS) Open(name string, opts ...OpenOption) (File, error) {
	f, err := fs.inner.Open(name, opts...)
	if f != nil {
		f = WithFd(f, enospcFile{
			fs:    fs,
			inner: f,
		})
	}
	return f, err
}

func (fs *enospcFS) OpenDir(name string) (File, error) {
	f, err := fs.inner.OpenDir(name)
	if f != nil {
		f = WithFd(f, enospcFile{
			fs:    fs,
			inner: f,
		})
	}
	return f, err
}

func (fs *enospcFS) Remove(name string) error {
	gen := fs.waitUntilReady()

	err := fs.inner.Remove(name)

	if err != nil && isENOSPC(err) {
		fs.handleENOSPC(gen)
		err = fs.inner.Remove(name)
	}
	return err
}

func (fs *enospcFS) RemoveAll(name string) error {
	gen := fs.waitUntilReady()

	err := fs.inner.RemoveAll(name)

	if err != nil && isENOSPC(err) {
		fs.handleENOSPC(gen)
		err = fs.inner.RemoveAll(name)
	}
	return err
}

func (fs *enospcFS) Rename(oldname, newname string) error {
	gen := fs.waitUntilReady()

	err := fs.inner.Rename(oldname, newname)

	if err != nil && isENOSPC(err) {
		fs.handleENOSPC(gen)
		err = fs.inner.Rename(oldname, newname)
	}
	return err
}

func (fs *enospcFS) ReuseForWrite(oldname, newname string) (File, error) {
	gen := fs.waitUntilReady()

	f, err := fs.inner.ReuseForWrite(oldname, newname)

	if err != nil && isENOSPC(err) {
		fs.handleENOSPC(gen)
		f, err = fs.inner.ReuseForWrite(oldname, newname)
	}

	if f != nil {
		f = WithFd(f, enospcFile{
			fs:    fs,
			inner: f,
		})
	}
	return f, err
}

func (fs *enospcFS) MkdirAll(dir string, perm os.FileMode) error {
	gen := fs.waitUntilReady()

	err := fs.inner.MkdirAll(dir, perm)

	if err != nil && isENOSPC(err) {
		fs.handleENOSPC(gen)
		err = fs.inner.MkdirAll(dir, perm)
	}
	return err
}

func (fs *enospcFS) Lock(name string) (io.Closer, error) {
	gen := fs.waitUntilReady()

	closer, err := fs.inner.Lock(name)

	if err != nil && isENOSPC(err) {
		fs.handleENOSPC(gen)
		closer, err = fs.inner.Lock(name)
	}
	return closer, err
}

func (fs *enospcFS) List(dir string) ([]string, error) {
	return fs.inner.List(dir)
}

func (fs *enospcFS) Stat(name string) (os.FileInfo, error) {
	return fs.inner.Stat(name)
}

func (fs *enospcFS) PathBase(path string) string {
	return fs.inner.PathBase(path)
}

func (fs *enospcFS) PathJoin(elem ...string) string {
	return fs.inner.PathJoin(elem...)
}

func (fs *enospcFS) PathDir(path string) string {
	return fs.inner.PathDir(path)
}

func (fs *enospcFS) GetDiskUsage(path string) (DiskUsage, error) {
	return fs.inner.GetDiskUsage(path)
}

type enospcFile struct {
	fs    *enospcFS
	inner File
}

func (f enospcFile) Close() error {
	return f.inner.Close()
}

func (f enospcFile) Read(p []byte) (n int, err error) {
	return f.inner.Read(p)
}

func (f enospcFile) ReadAt(p []byte, off int64) (n int, err error) {
	return f.inner.ReadAt(p, off)
}

func (f enospcFile) Write(p []byte) (n int, err error) {
	gen := f.fs.waitUntilReady()

	n, err = f.inner.Write(p)

	if err != nil && isENOSPC(err) {
		f.fs.handleENOSPC(gen)
		var n2 int
		n2, err = f.inner.Write(p[n:])
		n += n2
	}
	return n, err
}

func (f enospcFile) Stat() (os.FileInfo, error) {
	return f.inner.Stat()
}

func (f enospcFile) Sync() error {
	gen := f.fs.waitUntilReady()

	err := f.inner.Sync()

	if err != nil && isENOSPC(err) {
		f.fs.handleENOSPC(gen)

		// NB: It is NOT safe to retry the Sync. See the PostgreSQL
		// 'fsyncgate' discussion. A successful Sync after a failed one does
		// not provide any guarantees and (always?) loses the unsynced writes.
		// We need to bubble the error up and hope we weren't syncing a WAL or
		// MANIFEST, because we'll have no choice but to crash. Errors while
		// syncing an sstable will result in a failed flush/compaction, and
		// the relevant sstable(s) will be marked as obsolete and deleted.
		// See: https://lwn.net/Articles/752063/
	}
	return err
}

// Ensure that *enospcFS implements the FS interface.
var _ FS = (*enospcFS)(nil)

func isENOSPC(err error) bool {
	err = errors.UnwrapAll(err)
	e, ok := err.(syscall.Errno)
	return ok && e == syscall.ENOSPC
}