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 }