gcsbuf/writer.go

package buf

import (
	"bytes"
	"context"
	"fmt"
	"sync"
	"time"

	"cloud.google.com/go/storage"
)

const (
	maxComponentCount    = 32
	defaultFlushInterval = 5 * time.Second
	defaultFlushSize     = 4 * 1024 * 1024 // 4MB
)

var now = time.Now // var for testing

// WriterOptions configures a Writer.
type WriterOptions struct {
	// Client is the GCS client used for all operations.
	Client *storage.Client
	// Bucket is the GCS bucket name.
	Bucket string
	// Object is the GCS object name.
	Object string
	// FlushInterval is how often the background goroutine flushes buffered data. Default 5s.
	FlushInterval time.Duration
	// FlushSize is the buffer size threshold that triggers a synchronous flush on Write. Default 4MB.
	FlushSize int
	// OnOp, if set, is called after each GCS API operation completes.
	OnOp OnOp
}

type Writer struct {
	ctx    context.Context
	client *storage.Client
	bucket string
	object string

	flushInterval time.Duration
	flushSize     int

	mu      sync.Mutex
	buf     *bytes.Buffer
	created bool // whether the GCS object exists yet
	closed  bool
	err     error
	gen            int64 // expected generation of the GCS object
	componentCount int64 // tracked locally from compose results

	onOp OnOp

	stopCh chan struct{}
	doneCh chan struct{}
}

func NewWriter(ctx context.Context, opts WriterOptions) *Writer {
	interval := opts.FlushInterval
	if interval == 0 {
		interval = defaultFlushInterval
	}
	size := opts.FlushSize
	if size == 0 {
		size = defaultFlushSize
	}
	w := &Writer{
		ctx:           ctx,
		client:        opts.Client,
		bucket:        opts.Bucket,
		object:        opts.Object,
		flushInterval: interval,
		flushSize:     size,
		buf:           bytes.NewBuffer(nil),
		onOp:          opts.OnOp,
		stopCh:        make(chan struct{}),
		doneCh:        make(chan struct{}),
	}
	go w.backgroundFlush()
	return w
}

func (w *Writer) backgroundFlush() {
	defer close(w.doneCh)
	ticker := time.NewTicker(w.flushInterval)
	defer ticker.Stop()
	for {
		select {
		case <-w.stopCh:
			return
		case <-ticker.C:
			w.mu.Lock()
			if w.buf.Len() > 0 {
				w.flushLocked()
			}
			w.mu.Unlock()
		}
	}
}

func (w *Writer) Write(p []byte) (int, error) {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.closed {
		return 0, fmt.Errorf("write to closed writer")
	}
	if w.err != nil {
		return 0, w.err
	}
	n, err := w.buf.Write(p)
	if err != nil {
		return n, err
	}
	if w.buf.Len() >= w.flushSize {
		w.flushLocked()
	}
	if w.err != nil {
		return 0, w.err
	}
	return n, nil
}

// Flush flushes any buffered data to GCS.
func (w *Writer) Flush() error {
	w.mu.Lock()
	defer w.mu.Unlock()
	if w.buf.Len() > 0 {
		w.flushLocked()
	}
	return w.err
}

// flushLocked writes buffered data to GCS. Must be called with w.mu held.
func (w *Writer) flushLocked() {
	data := w.buf.Bytes()
	if len(data) == 0 {
		return
	}

	bh := w.client.Bucket(w.bucket)
	objh := bh.Object(w.object)

	if !w.created {
		// First write: create the object directly.
		// DoesNotExist ensures we don't overwrite an object created by someone else.
		attrs, err := retry(w.ctx, w.onOp, "upload", func() (*storage.ObjectAttrs, error) {
			gw := objh.If(storage.Conditions{DoesNotExist: true}).NewWriter(w.ctx)
			if _, err := gw.Write(data); err != nil {
				return nil, err
			}
			if err := gw.Close(); err != nil {
				return nil, err
			}
			return gw.Attrs(), nil
		})
		if err != nil {
			w.err = fmt.Errorf("writing initial object: %w", err)
			return
		}
		w.gen = attrs.Generation
		w.created = true
		w.buf.Reset()
		return
	}

	// Subsequent writes: write to temp, compose, delete temp.
	tmpname := fmt.Sprintf("%s-temp-%d", w.object, now().UnixNano())
	tmpobj := bh.Object(tmpname)

	if err := retryVoid(w.ctx, w.onOp, "upload", func() error {
		gw := tmpobj.If(storage.Conditions{DoesNotExist: true}).NewWriter(w.ctx)
		if _, err := gw.Write(data); err != nil {
			return err
		}
		return gw.Close()
	}); err != nil {
		w.err = fmt.Errorf("writing temp object: %w", err)
		return
	}

	// Check component count before composing.
	if w.componentCount >= maxComponentCount-1 {
		if err := w.resetComponentCount(bh, objh); err != nil {
			w.err = fmt.Errorf("resetting component count: %w", err)
			return
		}
	}

	// Compose existing + temp into existing.
	composeAttrs, err := retry(w.ctx, w.onOp, "compose", func() (*storage.ObjectAttrs, error) {
		return objh.If(storage.Conditions{GenerationMatch: w.gen}).ComposerFrom(objh, tmpobj).Run(w.ctx)
	})
	if err != nil {
		w.err = fmt.Errorf("composing objects: %w", err)
		return
	}
	w.gen = composeAttrs.Generation
	w.componentCount = composeAttrs.ComponentCount

	// Delete temp object.
	if err := retryVoid(w.ctx, w.onOp, "delete", func() error {
		return tmpobj.Delete(w.ctx)
	}); err != nil {
		w.err = fmt.Errorf("deleting temp object: %w", err)
		return
	}

	w.buf.Reset()
}

// resetComponentCount copies the object to a temp location and back to reset the component count.
func (w *Writer) resetComponentCount(bh *storage.BucketHandle, objh *storage.ObjectHandle) error {
	tmpname := fmt.Sprintf("%s-reset-%d", w.object, now().UnixNano())
	tmpobj := bh.Object(tmpname)

	// Copy to temp (this creates a non-composite object).
	// Source precondition ensures we're copying the generation we expect.
	if err := retryVoid(w.ctx, w.onOp, "rewrite", func() error {
		_, err := tmpobj.CopierFrom(objh.If(storage.Conditions{GenerationMatch: w.gen})).Run(w.ctx)
		return err
	}); err != nil {
		return fmt.Errorf("copying to reset temp: %w", err)
	}

	// Copy back.
	// Destination precondition ensures nobody modified the main object while we copied.
	attrs, err := retry(w.ctx, w.onOp, "rewrite", func() (*storage.ObjectAttrs, error) {
		return objh.If(storage.Conditions{GenerationMatch: w.gen}).CopierFrom(tmpobj).Run(w.ctx)
	})
	if err != nil {
		return fmt.Errorf("copying back from reset temp: %w", err)
	}
	w.gen = attrs.Generation

	// Delete the reset temp.
	if err := retryVoid(w.ctx, w.onOp, "delete", func() error {
		return tmpobj.Delete(w.ctx)
	}); err != nil {
		return fmt.Errorf("deleting reset temp: %w", err)
	}

	return nil
}

func (w *Writer) Close() error {
	w.mu.Lock()
	if w.closed {
		w.mu.Unlock()
		return fmt.Errorf("writer already closed")
	}
	w.closed = true
	w.mu.Unlock()

	// Stop background goroutine and wait for it.
	close(w.stopCh)
	<-w.doneCh

	// Final flush.
	w.mu.Lock()
	if w.buf.Len() > 0 {
		w.flushLocked()
	}
	err := w.err
	w.mu.Unlock()
	if err != nil {
		return err
	}

	// Set done metadata. Only if the object was created.
	if w.created {
		objh := w.client.Bucket(w.bucket).Object(w.object)
		if err := retryVoid(w.ctx, w.onOp, "patch", func() error {
			_, err := objh.If(storage.Conditions{GenerationMatch: w.gen}).Update(w.ctx, storage.ObjectAttrsToUpdate{
				Metadata: map[string]string{"done": "true"},
			})
			return err
		}); err != nil {
			return fmt.Errorf("setting done metadata: %w", err)
		}
	}

	return nil
}