gcsbuf/buf.go

package buf

import (
	"bytes"
	"context"
	"errors"
	"fmt"
	"io"
	"time"

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

const maxComponentCount = 32 // TODO verify

type Buffer struct {
	svc            *storage.Client
	bucket, object string
	buf            *bytes.Buffer // written; not yet flushed -- not read from directly

	readcur int
}

var _ interface {
	io.ReadCloser
	io.WriteCloser
} = (*Buffer)(nil)

type Options struct {
	Service        *storage.Client
	Bucket, Object string
	Cap            int
}

func NewBuffer(opts Options) *Buffer {
	return &Buffer{
		svc: opts.Service,
		buf: bytes.NewBuffer(nil),
	}
}

func (buf *Buffer) Write(b []byte) (int, error) { return -1, errors.New("unimplemented") }
func (buf *Buffer) Read(b []byte) (int, error)  { return -1, errors.New("unimplemented") }

func (buf *Buffer) Close() error {
	if err := buf.flush(); err != nil {
		return err
	}
	return buf.finalize()
}
func (buf *Buffer) Flush() error { return buf.flush() }

var now = time.Now // var for testing

// TODO: https://docs.cloud.google.com/storage/docs/request-preconditions to prevent races
func (buf *Buffer) flush() error {
	bh := buf.svc.Bucket(buf.bucket)

	// Insert a temp object.
	tmpname := fmt.Sprintf("%s-temp-%d", buf.object, now().UnixNano())
	tmpobj := bh.Object(tmpname)
	if _, err := io.Copy(tmpobj.NewWriter(context.TODO()), buf.buf); err != nil {
		return fmt.Errorf("writing temp object: %w", err)
	}

	// Compose it into the running object.
	attrs, err := buf.svc.Bucket(buf.bucket).Object(buf.object).ComposerFrom(tmpobj).Run(context.TODO())
	if err != nil {
		return fmt.Errorf("composing temp object: %w", err)
	}

	// If reaching the max composite count, copy and rename to reset composite count.
	if attrs.ComponentCount >= maxComponentCount {
		// Copy object and move back
		tmpname = fmt.Sprintf("%s-temp-%d", buf.object, now().UnixNano())
		objh := bh.Object(buf.object)
		if _, err := bh.Object(tmpname).CopierFrom(objh).Run(context.TODO()); err != nil {
			return fmt.Errorf("copying temp object: %w", err)
		}
		if _, err := bh.Object(tmpname).Move(context.TODO(), storage.MoveObjectDestination{Object: buf.object}); err != nil {
			return fmt.Errorf("renaming temp object: %w", err)
		}
	}
	return nil
}

func (buf *Buffer) finalize() error {
	if _, err := buf.svc.Bucket(buf.bucket).Object(buf.object).Update(context.TODO(), storage.ObjectAttrsToUpdate{
		Metadata: map[string]string{"done": "true"},
	}); err != nil {
		return fmt.Errorf("finalizing object: %w", err)
	}
	return nil
}