add test binary, observability

Signed-off-by: Jason Hall <imjasonh@gmail.com>

README.md

@@ -33,6 +33,37 @@ _, _ = io.Copy(os.Stdout, r) // reads until the object is finalized
 
 In the example above, if the writer is still writing data, `io.Copy` will block until the writer completes, and receive new content each time it flushes.
 
+If the writer has already written and finalized its data, `io.Copy` simply streams all of the data in one request.
+
+### Rate limits
+
+GCS enforces a [1 mutation per second per object](https://cloud.google.com/storage/quotas) rate limit. Each flush cycle performs one compose on the main object (temp object operations target separate objects), so the default 5-second flush interval stays well within this limit. If you use manual `Flush()` calls or a low `FlushInterval` or write data faster than your `FlushSize` per second, you may hit rate limits; all GCS operations are retried automatically with exponential backoff, but this will slow down your writes.
+
+### Observability
+
+Both `WriterOptions` and `ReaderOptions` accept an optional `OnOp` callback, called after each GCS API operation with the operation name, duration, and error (nil on success). Operations: `upload`, `attrs`, `compose`, `delete`, `rewrite`, `patch`, `read`.
+
+```go
+w := gcsbuf.NewWriter(ctx, gcsbuf.WriterOptions{
+    Client: client,
+    Bucket: "my-bucket",
+    Object: "my-object",
+    OnOp: func(op string, d time.Duration, err error) {
+        slog.Info("gcs", "op", op, "dur", d, "err", err)
+    },
+})
+```
+
+You can use this to export Prometheus metrics for each operation.
+
+### Testing
+
+You can test this with `./cmd/gcscat`, which reads from `stdin`, streamily writes it to GCS, and simultaneously reads it back out, and logs operations.
+
+```
+echo "hello world" | go run ./cmd/gcscat <bucket> <object>
+```
+
 ### Background
 
 I originally used a ~hack~ scheme like this when creating [Google Cloud Build](https://cloud.google.com/build) to efficiently write logs to durable storage, and to allow clients like `gcloud` to simultaneously stream the logs objects to users until build completion. Logs were separately written to Google Cloud Logging, which didn't support a fast logs tailing API at the time.

cmd/gcscat/main.go

@@ -0,0 +1,72 @@
+package main
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"os"
+	"sync"
+	"time"
+
+	"cloud.google.com/go/storage"
+	buf "github.com/imjasonh/gcsbuf"
+)
+
+func main() {
+	if len(os.Args) != 3 {
+		fmt.Fprintf(os.Stderr, "usage: gcscat <bucket> <object>\n")
+		os.Exit(1)
+	}
+	bucket, object := os.Args[1], os.Args[2]
+
+	ctx := context.Background()
+	client, err := storage.NewClient(ctx)
+	if err != nil {
+		slog.Error("creating storage client", "err", err)
+		os.Exit(1)
+	}
+	defer client.Close()
+
+	onOp := buf.OnOp(func(op string, d time.Duration, err error) {
+		if err != nil {
+			slog.Error("gcs", "op", op, "dur", d, "err", err)
+		} else {
+			slog.Info("gcs", "op", op, "dur", d)
+		}
+	})
+
+	w := buf.NewWriter(ctx, buf.WriterOptions{
+		Client: client,
+		Bucket: bucket,
+		Object: object,
+		OnOp:   onOp,
+	})
+
+	r := buf.NewReader(ctx, buf.ReaderOptions{
+		Client: client,
+		Bucket: bucket,
+		Object: object,
+		OnOp:   onOp,
+	})
+
+	// Read concurrently, printing to stdout.
+	var wg sync.WaitGroup
+	wg.Go(func() {
+		if _, err := io.Copy(os.Stdout, r); err != nil {
+			slog.Error("read", "err", err)
+		}
+	})
+
+	// Write stdin to GCS.
+	if _, err := io.Copy(w, os.Stdin); err != nil {
+		slog.Error("write", "err", err)
+		os.Exit(1)
+	}
+	if err := w.Close(); err != nil {
+		slog.Error("close writer", "err", err)
+		os.Exit(1)
+	}
+
+	wg.Wait()
+}

observe.go

@@ -0,0 +1,14 @@
+package buf
+
+import "time"
+
+// OnOp is called after each GCS API operation completes.
+// op identifies the operation (e.g. "upload", "attrs", "compose", "delete", "rewrite", "patch", "read").
+// d is the wall-clock duration. err is nil on success.
+type OnOp func(op string, d time.Duration, err error)
+
+func (f OnOp) observe(op string, d time.Duration, err error) {
+	if f != nil {
+		f(op, d, err)
+	}
+}

reader.go

@@ -24,6 +24,8 @@ type ReaderOptions struct {
 	Object string
 	// PollInterval is how often the reader polls for new data or object creation. Default 1s.
 	PollInterval time.Duration
+	// OnOp, if set, is called after each GCS API operation completes.
+	OnOp OnOp
 }
 
 type Reader struct {
@@ -32,6 +34,7 @@ type Reader struct {
 	bucket       string
 	object       string
 	pollInterval time.Duration
+	onOp         OnOp
 
 	cursor int64
 	done   bool
@@ -48,6 +51,7 @@ func NewReader(ctx context.Context, opts ReaderOptions) *Reader {
 		bucket:       opts.Bucket,
 		object:       opts.Object,
 		pollInterval: interval,
+		onOp:         opts.OnOp,
 	}
 }
 
@@ -55,7 +59,9 @@ func (r *Reader) Read(p []byte) (int, error) {
 	for {
 		objh := r.client.Bucket(r.bucket).Object(r.object)
 
-		attrs, err := objh.Attrs(r.ctx)
+		attrs, err := retry(r.ctx, r.onOp, "attrs", func() (*storage.ObjectAttrs, error) {
+			return objh.Attrs(r.ctx)
+		})
 		if err != nil {
 			if isNotFound(err) {
 				// Object doesn't exist yet, poll.
@@ -87,13 +93,19 @@ func (r *Reader) Read(p []byte) (int, error) {
 			length = available
 		}
 
-		rc, err := objh.NewRangeReader(r.ctx, r.cursor, length)
+		n, err := retry(r.ctx, r.onOp, "read", func() (int, error) {
+			rc, err := objh.NewRangeReader(r.ctx, r.cursor, length)
+			if err != nil {
+				return 0, err
+			}
+			n, err := io.ReadFull(rc, p[:length])
+			rc.Close()
+			if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) {
+				return n, err
+			}
+			return n, nil
+		})
 		if err != nil {
-			return 0, fmt.Errorf("creating range reader: %w", err)
-		}
-		n, err := io.ReadFull(rc, p[:length])
-		rc.Close()
-		if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) {
 			return n, fmt.Errorf("reading object data: %w", err)
 		}
 

retry.go

@@ -0,0 +1,50 @@
+package buf
+
+import (
+	"context"
+	"errors"
+	"math/rand/v2"
+	"net/http"
+	"time"
+
+	"google.golang.org/api/googleapi"
+)
+
+const maxRetries = 5
+
+func isRateLimited(err error) bool {
+	var apiErr *googleapi.Error
+	return errors.As(err, &apiErr) && apiErr.Code == http.StatusTooManyRequests
+}
+
+// retry retries fn on rate-limit (429) errors with exponential backoff and jitter.
+// OnOp is fired after each attempt, so callers see every rate-limit hit.
+func retry[T any](ctx context.Context, onOp OnOp, op string, fn func() (T, error)) (T, error) {
+	for attempt := range maxRetries + 1 {
+		start := time.Now()
+		v, err := fn()
+		onOp.observe(op, time.Since(start), err)
+		if err == nil || !isRateLimited(err) {
+			return v, err
+		}
+		if attempt == maxRetries {
+			return v, err
+		}
+		backoff := min(time.Duration(1<<attempt)*100*time.Millisecond, 5*time.Second)
+		jitter := time.Duration(rand.Int64N(int64(backoff) / 2))
+		select {
+		case <-time.After(backoff + jitter):
+		case <-ctx.Done():
+			var zero T
+			return zero, ctx.Err()
+		}
+	}
+	panic("unreachable")
+}
+
+func retryVoid(ctx context.Context, onOp OnOp, op string, fn func() error) error {
+	_, err := retry(ctx, onOp, op, func() (struct{}, error) {
+		return struct{}{}, fn()
+	})
+	return err
+}

writer.go

@@ -30,6 +30,8 @@ type WriterOptions struct {
 	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 {
@@ -46,7 +48,10 @@ type Writer struct {
 	created bool // whether the GCS object exists yet
 	closed  bool
 	err     error
-	gen     int64 // expected generation of the GCS object
+	gen            int64 // expected generation of the GCS object
+	componentCount int64 // tracked locally from compose results
+
+	onOp OnOp
 
 	stopCh chan struct{}
 	doneCh chan struct{}
@@ -69,6 +74,7 @@ func NewWriter(ctx context.Context, opts WriterOptions) *Writer {
 		flushInterval: interval,
 		flushSize:     size,
 		buf:           bytes.NewBuffer(nil),
+		onOp:          opts.OnOp,
 		stopCh:        make(chan struct{}),
 		doneCh:        make(chan struct{}),
 	}
@@ -139,16 +145,21 @@ func (w *Writer) flushLocked() {
 	if !w.created {
 		// First write: create the object directly.
 		// DoesNotExist ensures we don't overwrite an object created by someone else.
-		gw := objh.If(storage.Conditions{DoesNotExist: true}).NewWriter(w.ctx)
-		if _, err := gw.Write(data); err != nil {
+		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
 		}
-		if err := gw.Close(); err != nil {
-			w.err = fmt.Errorf("closing initial object writer: %w", err)
-			return
-		}
-		w.gen = gw.Attrs().Generation
+		w.gen = attrs.Generation
 		w.created = true
 		w.buf.Reset()
 		return
@@ -158,24 +169,19 @@ func (w *Writer) flushLocked() {
 	tmpname := fmt.Sprintf("%s-temp-%d", w.object, now().UnixNano())
 	tmpobj := bh.Object(tmpname)
 
-	gw := tmpobj.If(storage.Conditions{DoesNotExist: true}).NewWriter(w.ctx)
-	if _, err := gw.Write(data); err != nil {
+	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
 	}
-	if err := gw.Close(); err != nil {
-		w.err = fmt.Errorf("closing temp object writer: %w", err)
-		return
-	}
 
 	// Check component count before composing.
-	attrs, err := objh.If(storage.Conditions{GenerationMatch: w.gen}).Attrs(w.ctx)
-	if err != nil {
-		w.err = fmt.Errorf("getting object attrs: %w", err)
-		return
-	}
-	if attrs.ComponentCount >= maxComponentCount-1 {
-		// Reset component count: copy to temp2, copy back, delete temp2.
+	if w.componentCount >= maxComponentCount-1 {
 		if err := w.resetComponentCount(bh, objh); err != nil {
 			w.err = fmt.Errorf("resetting component count: %w", err)
 			return
@@ -183,16 +189,20 @@ func (w *Writer) flushLocked() {
 	}
 
 	// Compose existing + temp into existing.
-	composer := objh.If(storage.Conditions{GenerationMatch: w.gen}).ComposerFrom(objh, tmpobj)
-	composeAttrs, err := composer.Run(w.ctx)
+	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 := tmpobj.Delete(w.ctx); err != nil {
+	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
 	}
@@ -207,20 +217,27 @@ func (w *Writer) resetComponentCount(bh *storage.BucketHandle, objh *storage.Obj
 
 	// Copy to temp (this creates a non-composite object).
 	// Source precondition ensures we're copying the generation we expect.
-	if _, err := tmpobj.CopierFrom(objh.If(storage.Conditions{GenerationMatch: w.gen})).Run(w.ctx); err != nil {
+	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 := objh.If(storage.Conditions{GenerationMatch: w.gen}).CopierFrom(tmpobj).Run(w.ctx)
+	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 := tmpobj.Delete(w.ctx); err != nil {
+	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)
 	}
 
@@ -254,8 +271,11 @@ func (w *Writer) Close() error {
 	// Set done metadata. Only if the object was created.
 	if w.created {
 		objh := w.client.Bucket(w.bucket).Object(w.object)
-		if _, err := objh.If(storage.Conditions{GenerationMatch: w.gen}).Update(w.ctx, storage.ObjectAttrsToUpdate{
-			Metadata: map[string]string{"done": "true"},
+		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)
 		}