go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/common/gcloud/gs/writer.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package gs
    16  
    17  import (
    18  	"io"
    19  	"time"
    20  
    21  	gs "cloud.google.com/go/storage"
    22  	log "go.chromium.org/luci/common/logging"
    23  	"go.chromium.org/luci/common/retry"
    24  	"go.chromium.org/luci/common/retry/transient"
    25  )
    26  
    27  // Writer is an augmented io.WriteCloser instance.
    28  type Writer interface {
    29  	io.WriteCloser
    30  
    31  	// Count returns the number of bytes written by the object.
    32  	Count() int64
    33  }
    34  
    35  type prodWriter struct {
    36  	// Writer is the active Writer instance. It will be nil until the first Write
    37  	// invocation.
    38  	writer *gs.Writer
    39  
    40  	client  *prodClient
    41  	bucket  string
    42  	relpath string
    43  	count   int64
    44  }
    45  
    46  var _ Writer = (*prodWriter)(nil)
    47  
    48  // Write writes data with exponential backoff/retry.
    49  func (w *prodWriter) Write(d []byte) (a int, err error) {
    50  	if w.writer == nil {
    51  		w.writer = w.client.baseClient.Bucket(w.bucket).Object(w.relpath).NewWriter(w.client.ctx)
    52  	}
    53  
    54  	err = retry.Retry(w.client.ctx, transient.Only(retry.Default), func() (ierr error) {
    55  		a, ierr = w.writer.Write(d)
    56  
    57  		// Assume all Write errors are transient.
    58  		ierr = transient.Tag.Apply(ierr)
    59  		return
    60  	}, func(err error, d time.Duration) {
    61  		log.Fields{
    62  			log.ErrorKey: err,
    63  			"delay":      d,
    64  			"bucket":     w.bucket,
    65  			"path":       w.relpath,
    66  		}.Warningf(w.client.ctx, "Transient error on GS write. Retrying...")
    67  	})
    68  
    69  	w.count += int64(a)
    70  	return
    71  }
    72  
    73  func (w *prodWriter) Close() error {
    74  	if w.writer == nil {
    75  		return nil
    76  	}
    77  
    78  	return retry.Retry(w.client.ctx, transient.Only(retry.Default),
    79  		w.writer.Close,
    80  		func(err error, d time.Duration) {
    81  			log.Fields{
    82  				log.ErrorKey: err,
    83  				"delay":      d,
    84  				"bucket":     w.bucket,
    85  				"path":       w.relpath,
    86  			}.Warningf(w.client.ctx, "Transient error closing GS Writer. Retrying...")
    87  		})
    88  }
    89  
    90  func (w *prodWriter) Count() int64 {
    91  	return w.count
    92  }