// Copyright 2014 Google LLC

//

// Licensed under the Apache License, Version 2.0 (the "License");

// you may not use this file except in compliance with the License.

// You may obtain a copy of the License at

//

// http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing, software

// distributed under the License is distributed on an "AS IS" BASIS,

// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

// See the License for the specific language governing permissions and

// limitations under the License.

package storage

import (

"context"

"errors"

"fmt"

"io"

"log"

"sync"

"sync/atomic"

"time"

"unicode/utf8"

"go.opentelemetry.io/otel/attribute"

"go.opentelemetry.io/otel/metric"

)

// Interface internalWriter wraps low-level implementations which may vary

// across client types.

type internalWriter interface {

io.WriteCloser

Flush() (int64, error)

// CloseWithError terminates the write operation and sets its status.

// Note that CloseWithError always returns nil.

CloseWithError(error) error

// Set final object checksum for appendable objects after write

// and before finalization.

setAppendFinalCRC32C(sendAppendFinalCRC32C bool, c uint32)

}

// A Writer writes a Cloud Storage object.

type Writer struct {

// ObjectAttrs are optional attributes to set on the object. Any attributes

// must be initialized before the first Write call. Nil or zero-valued

// attributes are ignored.

ObjectAttrs

// SendCRC32C specifies whether to transmit a CRC32C checksum. When this is

// true and the Writer's CRC32C field is set, that checksum is sent to GCS.

// If the data written does not match the checksum, the write is rejected.

// It is necessary to set this field to true in addition to setting the

// Writer's CRC32C field because zero is a valid CRC.

//

// By default, the client automatically calculates and sends checksums.

// When using gRPC, checksums are sent for both individual chunks and the full object.

// When using JSON, checksums are sent only for the full object.

// However, a user-provided checksum takes precedence over the auto-calculated checksum

// for the full object.

//

// Note: SendCRC32C must be set before the first call to Writer.Write().

SendCRC32C bool

// DisableAutoChecksum disables automatic CRC32C checksum calculation and

// validation in the Writer. By default, the Writer automatically performs

// checksum validation. Setting this to true disables this behavior.

//

// Disabling automatic checksumming does not prevent a user-provided checksum

// from being sent. If SendCRC32C is true and the Writer's CRC32C field is

// populated, that checksum will still be sent to GCS for validation.

//

// For single-shot JSON uploads, a mismatch in the auto-calculated checksum returns

// an error but may leave data on the server. This issue does not apply when

// user-provided checksum is used. Callers relying on auto-checksum should handle the

// error by removing the object or restoring a previous version.

//

// Automatic CRC32C checksum calculation introduces increased CPU overhead

// because of checksum computation in writes. Use this field to disable

// it if needed.

//

// Note: DisableAutoChecksum must be set before the first call to

// Writer.Write(). Automatic checksumming is not enabled for full object

// checksums for unfinalized writes to appendable objects in gRPC.

DisableAutoChecksum bool

// ChunkSize controls the maximum number of bytes of the object that the

// Writer will attempt to send to the server in a single request. Objects

// smaller than the size will be sent in a single request, while larger

// objects will be split over multiple requests. The value will be rounded up

// to the nearest multiple of 256K. The default ChunkSize is 16MiB.

//

// Each Writer will internally allocate a buffer of size ChunkSize. This is

// used to buffer input data and allow for the input to be sent again if a

// request must be retried.

//

// If you upload small objects (< 16MiB), you should set ChunkSize

// to a value slightly larger than the objects' sizes to avoid memory bloat.

// This is especially important if you are uploading many small objects

// concurrently. See

// https://cloud.google.com/storage/docs/json_api/v1/how-tos/upload#size

// for more information about performance trade-offs related to ChunkSize.

//

// If ChunkSize is set to zero, chunking will be disabled and the object will

// be uploaded in a single request without the use of a buffer. This will

// further reduce memory used during uploads, but will also prevent the writer

// from retrying in case of a transient error from the server or resuming an

// upload that fails midway through, since the buffer is required in order to

// retry the failed request.

//

// ChunkSize must be set before the first Write call.

ChunkSize int

// ChunkRetryDeadline sets a per-chunk retry deadline for uploads.

//

// For uploads of larger files, the Writer will attempt to retry if the

// request to upload a particular chunk fails with a transient error.

// If a single chunk has been attempting to upload for longer than this

// deadline and the request fails, it will no longer be retried, and the

// error will be returned to the caller. This deadline measures the total

// time spent queuing, transmitting, and retrying a single chunk.

//

// For HTTP clients, a chunk is defined by the ChunkSize, so this deadline

// is only applicable to files large enough to require a multi-chunk upload.

// Users may also want to pick a longer deadline if they are using larger

// values for ChunkSize.

//

// For gRPC clients, data is streamed internally in 2 MiB quantums,

// and this deadline applies to each individual quantum. Therefore, it applies

// to all uploads, and the deadline does not need to scale with the ChunkSize.

//

// The default value is 32s. Users may want to pick a longer deadline if they

// expect to have a slow or unreliable internet connection.

//

// To set a deadline on the entire upload, use context timeout or

// cancellation.

ChunkRetryDeadline time.Duration

// ChunkTransferTimeout sets a per-chunk request timeout for resumable uploads.

//

// For resumable uploads, the Writer will terminate the request and attempt

// a retry if the request to upload a particular chunk stalls for longer than

// this duration. Retries may continue until the ChunkRetryDeadline is reached.

//

// ChunkTransferTimeout is not applicable to uploads made using a gRPC client.

//

// The default value is no timeout.

ChunkTransferTimeout time.Duration

// ForceEmptyContentType is an optional parameter that is used to disable

// auto-detection of Content-Type. By default, if a blank Content-Type

// is provided, then gax.DetermineContentType is called to sniff the type.

ForceEmptyContentType bool

// Append is a parameter to indicate whether the writer should use appendable

// object semantics for the new object generation. Appendable objects are

// visible on the first Write() call, and can be appended to until they are

// finalized. If Writer.FinalizeOnClose is set to true, the object is finalized

// when Writer.Close() is called; otherwise, the object is left unfinalized

// and can be appended to later.

//

// Defaults to false unless the [WithAppendableUploads] option was set.

//

// Append is only supported for gRPC.

Append bool

// FinalizeOnClose indicates whether the Writer should finalize an object when

// closing the write stream. This only applies to Writers where Append is

// true, since append semantics allow a prefix of the object to be durable and

// readable. By default, objects written with Append semantics will not be

// finalized, which means they can be appended to later. If Append is set

// to false, this parameter will be ignored; non-appendable objects will

// always be finalized when Writer.Close returns without error.

FinalizeOnClose bool

// ProgressFunc can be used to monitor the progress of a large write

// operation. If ProgressFunc is not nil and writing requires multiple

// calls to the underlying service (see

// https://cloud.google.com/storage/docs/json_api/v1/how-tos/resumable-upload),

// then ProgressFunc will be invoked after each call with the number of bytes of

// content copied so far.

//

// For parallel uploads, progress is reported when each part is successfully uploaded.

// Therefore, the progress may be delayed relative to the standard upload,

// and jump in increments of the part size used (e.g. 16MiB).

//

// ProgressFunc should return quickly without blocking.

ProgressFunc func(int64)

// EnableParallelUpload enables the parallel upload feature.

// This feature splits a large object into multiple parts and uploads them in

// parallel. Supported exclusively for gRPC clients. If used with a JSON

// client, the configuration is ignored and a standard upload is performed.

//

// Parallel uploads can yield higher throughput when uploading large objects,

// but there are several considerations and trade-offs. Please refer to

// the [Parallel Uploads] section in the package documentation for details.

//

// **Note:** This feature is currently experimental and its API surface may change

// in future releases. It is not yet recommended for production use.

EnableParallelUpload bool

// ParallelUploadConfig holds configuration for Parallel Uploads.

// This only takes effect if EnableParallelUpload is true. Supported

// exclusively for gRPC clients. If used with a JSON client, the configuration

// is ignored and a standard upload is performed.

//

// **Note:** This feature is currently experimental and its API surface may change

// in future releases. It is not yet recommended for production use.

ParallelUploadConfig ParallelUploadConfig

// SendAppendFinalCRC32C indicates that AppendFinalCRC32C should be sent as the

// full-object checksum when finalizing an appendable object.

//

// This field is only supported for gRPC clients and appendable objects.

SendAppendFinalCRC32C bool

// AppendFinalCRC32C is the expected full-object CRC32C checksum to be validated

// by the server when finalizing an appendable object.

//

// This field must be set on the Writer, along with SendAppendFinalCRC32C = true,

// before calling Close(). It allows callers to defer providing the checksum

// until the final write is complete.

//

// If SendAppendFinalCRC32C is false, checksum validation for the full object will

// fall back to using ObjectAttrs.CRC32C if Writer.SendCRC32C is true. If both are

// configured, AppendFinalCRC32C takes precedence.

//

// This field is ignored if Writer.Append is false or if using the JSON API.

// Chunk-level validation will still be performed for all chunks during upload if

// Writer.DisableAutoChecksum is false.

AppendFinalCRC32C uint32

ctx context.Context

o *ObjectHandle

pcu *pcuState

opened bool

closed bool

iw internalWriter

donec chan struct{} // closed after err and obj are set.

obj *ObjectAttrs

mu sync.Mutex

err error

setTakeoverOffset func(int64)

// bytesWritten is the cumulative bytes written for request size metric.

bytesWritten int64

}

func (w *Writer) wrapWriteError(n int, err error) (int, error) {

if err == nil {

return n, nil

}

w.mu.Lock()

werr := w.err

w.mu.Unlock()

// Preserve existing functionality that when context is canceled, Write will return

// context.Canceled instead of "io: read/write on closed pipe". This hides the

// pipe implementation detail from users and makes Write seem as though it's an RPC.

if errors.Is(werr, context.Canceled) || errors.Is(werr, context.DeadlineExceeded) {

return n, werr

}

return n, err

}

func (w *Writer) isGRPCClient() bool {

tc := w.o.c.tc

if mc, ok := tc.(*metricsStorageClient); ok {

tc = mc.storageClient

}

_, ok := tc.(*grpcStorageClient)

return ok

}

func (w *Writer) getOrInitPCU() (*pcuState, error) {

w.mu.Lock()

defer w.mu.Unlock()

if w.pcu == nil {

if !w.EnableParallelUpload {

return nil, nil

}

if !w.isGRPCClient() {

// PCU only supported for gRPC.

// Nullify the config and proceed with standard upload.

log.Printf("storage: ParallelUploadConfig is ignored because Parallel Uploads are only supported for gRPC clients. Proceeding with standard upload.")

w.EnableParallelUpload = false

return nil, nil

}

if err := w.initPCU(w.ctx); err != nil {

return nil, err

}

}

return w.pcu, nil

}

// Write appends to w. It implements the io.Writer interface.

//

// Since writes happen asynchronously, Write may return a nil

// error even though the write failed (or will fail). Always

// use the error returned from Writer.Close to determine if

// the upload was successful.

//

// Writes will be retried on transient errors from the server, unless

// Writer.ChunkSize has been set to zero.

func (w *Writer) Write(p []byte) (int, error) {

w.mu.Lock()

werr, closed, pcu := w.err, w.closed, w.pcu

w.mu.Unlock()

if werr != nil {

return 0, werr

}

if closed {

return 0, fmt.Errorf("storage: Writer is closed")

}

var n int

var err error

if pcu != nil {

n, err = pcu.write(p)

} else {

if !w.opened {

// First time initialization: freeze the configuration to either PCU or standard.

if w.EnableParallelUpload {

if pcu, err = w.getOrInitPCU(); err != nil {

return 0, err

}

}

if pcu == nil {

if err = w.openWriter(); err != nil {

return 0, err

}

}

}

if pcu != nil {

n, err = pcu.write(p)

} else {

n, err = w.iw.Write(p)

}

}

if n > 0 {

atomic.AddInt64(&w.bytesWritten, int64(n))

}

return w.wrapWriteError(n, err)

}

// Flush syncs all bytes currently in the Writer's buffer to Cloud Storage.

// It returns the offset of bytes that have been currently synced to

// Cloud Storage and an error.

//

// If Flush is never called, Writer will sync data automatically every

// [Writer.ChunkSize] bytes and on [Writer.Close].

//

// [Writer.ProgressFunc] will be called on Flush if present.

//

// Do not call Flush concurrently with Write or Close. A single Writer is not

// safe for unsynchronized use across threads.

//

// Note that calling Flush very early (before 512 bytes) may interfere with

// automatic content sniffing in the Writer.

//

// Flush is supported only on gRPC clients where [Writer.Append] is set

// to true.

func (w *Writer) Flush() (int64, error) {

// Return error if Append is not true.

if !w.Append {

return 0, errors.New("storage: Flush not supported unless client uses gRPC and Append is set to true")

}

if w.closed {

return 0, errors.New("storage: Flush called on closed Writer")

}

// Return error if already in error state.

w.mu.Lock()

werr := w.err

w.mu.Unlock()

if werr != nil {

return 0, werr

}

// If Flush called before any bytes written, it should start the upload

// at zero bytes. This will make the object visible with zero length data.

if !w.opened {

if err := w.openWriter(); err != nil {

return 0, err

}

}

return w.iw.Flush()

}

// Close completes the write operation and flushes any buffered data.

// If Close doesn't return an error, metadata about the written object

// can be retrieved by calling Attrs.

func (w *Writer) Close() error {

w.mu.Lock()

closed, werr, pcu := w.closed, w.err, w.pcu

w.mu.Unlock()

if closed {

return werr

}

if pcu != nil || (!w.opened && w.EnableParallelUpload) {

var err error

if pcu, err = w.getOrInitPCU(); err != nil {

return w.markClosed(err)

}

if pcu != nil {

return w.markClosed(pcu.close())

}

}

if !w.opened {

if err := w.openWriter(); err != nil {

return w.markClosed(err)

}

}

if w.Append {

w.iw.setAppendFinalCRC32C(w.SendAppendFinalCRC32C, w.AppendFinalCRC32C)

}

if err := w.iw.Close(); err != nil {

return w.markClosed(err)

}

<-w.donec

return w.markClosed(nil)

}

// Abort is unimplemented and always returns an error.

//

// This is experimental and its signature can change in the future.

func (w *Writer) Abort() error {

return errMethodNotSupported

}

// markClosed marks the Writer as closed, records any closing error on Writer.err,

// and records request body size metrics and trace span completion.

func (w *Writer) markClosed(err error) error {

w.mu.Lock()

w.closed = true

if w.err == nil && err != nil {

w.err = err

}

closingErr := w.err

total := atomic.LoadInt64(&w.bytesWritten)

w.mu.Unlock()

if state := metricsStateFromContext(w.ctx); state != nil {

if state.metrics != nil && total > 0 {

state.metrics.requestBodySize.Record(w.ctx, total, metric.WithAttributes(attribute.String("rpc.system.name", state.getSystemName()), attribute.String("rpc.method", "WriteObject"), attribute.String("server.address", stripPort(state.getTarget()))))

}

if state.record != nil {

state.record(closingErr)

}

}

endSpan(w.ctx, closingErr)

return closingErr

}

func (w *Writer) openWriter() (err error) {

if err := w.validateWriteAttrs(); err != nil {

return err

}

if w.o.gen != defaultGen && !w.Append {

return fmt.Errorf("storage: generation supported on Writer for appendable objects only, got %v", w.o.gen)

}

isIdempotent := w.o.conds != nil && (w.o.conds.GenerationMatch >= 0 || w.o.conds.DoesNotExist)

// Append operations that takeover a specific generation are idempotent.

isIdempotent = isIdempotent || w.Append && w.o.gen > 0

opts := makeStorageOpts(isIdempotent, w.o.retry, w.o.userProject)

params := &openWriterParams{

ctx: w.ctx,

chunkSize: w.ChunkSize,

chunkRetryDeadline: w.ChunkRetryDeadline,

chunkTransferTimeout: w.ChunkTransferTimeout,

bucket: w.o.bucket,

attrs: &w.ObjectAttrs,

conds: w.o.conds,

appendGen: w.o.gen,

encryptionKey: w.o.encryptionKey,

sendCRC32C: w.SendCRC32C,

disableAutoChecksum: w.DisableAutoChecksum,

append: w.Append,

finalizeOnClose: w.FinalizeOnClose,

donec: w.donec,

setError: w.error,

progress: w.progress,

setObj: func(o *ObjectAttrs) { w.obj = o },

setSize: func(n int64) {

if w.obj != nil {

w.obj.Size = n

}

},

setTakeoverOffset: w.setTakeoverOffset,

forceEmptyContentType: w.ForceEmptyContentType,

}

if err := w.ctx.Err(); err != nil {

return err // short-circuit

}

w.iw, err = w.o.c.tc.OpenWriter(params, opts...)

if err != nil {

return err

}

w.ctx = params.ctx

w.opened = true

go w.monitorCancel()

return nil

}

// monitorCancel is intended to be used as a background goroutine. It monitors the

// context, and when it observes that the context has been canceled, it manually

// closes things that do not take a context.

func (w *Writer) monitorCancel() {

select {

case <-w.ctx.Done():

w.mu.Lock()

werr := w.ctx.Err()

w.err = werr

w.mu.Unlock()

w.CloseWithError(werr)

case <-w.donec:

}

}

// CloseWithError aborts the write operation with the provided error.

// CloseWithError always returns nil.

//

// Deprecated: cancel the context passed to NewWriter instead.

func (w *Writer) CloseWithError(err error) error {

if !w.opened {

return nil

}

return w.iw.CloseWithError(err)

}

// Attrs returns metadata about a successfully-written object.

// It's only valid to call it after Close returns nil.

func (w *Writer) Attrs() *ObjectAttrs {

return w.obj

}

func (w *Writer) validateWriteAttrs() error {

attrs := w.ObjectAttrs

// Check the developer didn't change the object Name (this is unfortunate, but

// we don't want to store an object under the wrong name).

if attrs.Name != w.o.object {

return fmt.Errorf("storage: Writer.Name %q does not match object name %q", attrs.Name, w.o.object)

}

if !utf8.ValidString(attrs.Name) {

return fmt.Errorf("storage: object name %q is not valid UTF-8", attrs.Name)

}

if attrs.KMSKeyName != "" && w.o.encryptionKey != nil {

return errors.New("storage: cannot use KMSKeyName with a customer-supplied encryption key")

}

if w.ChunkSize < 0 {

return errors.New("storage: Writer.ChunkSize must be non-negative")

}

return nil

}

// progress is a convenience wrapper that reports write progress to the Writer

// ProgressFunc if it is set.

func (w *Writer) progress(p int64) {

if w.ProgressFunc != nil {

w.ProgressFunc(p)

}

}

// error acquires the Writer's lock, sets the Writer's err to the given error,

// then relinquishes the lock.

func (w *Writer) error(err error) {

w.mu.Lock()

w.err = err

w.mu.Unlock()

}