vendor/cloud.google.com/go/storage/writer.go

   1 // Copyright 2014 Google LLC
   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 storage
  16
  17 import (
  18         "context"
  19         "encoding/base64"
  20         "errors"
  21         "fmt"
  22         "io"
  23         "sync"
  24         "unicode/utf8"
  25
  26         "google.golang.org/api/googleapi"
  27         raw "google.golang.org/api/storage/v1"
  28 )
  29
  30 // A Writer writes a Cloud Storage object.
  31 type Writer struct {
  32         // ObjectAttrs are optional attributes to set on the object. Any attributes
  33         // must be initialized before the first Write call. Nil or zero-valued
  34         // attributes are ignored.
  35         ObjectAttrs
  36
  37         // SendCRC specifies whether to transmit a CRC32C field. It should be set
  38         // to true in addition to setting the Writer's CRC32C field, because zero
  39         // is a valid CRC and normally a zero would not be transmitted.
  40         // If a CRC32C is sent, and the data written does not match the checksum,
  41         // the write will be rejected.
  42         SendCRC32C bool
  43
  44         // ChunkSize controls the maximum number of bytes of the object that the
  45         // Writer will attempt to send to the server in a single request. Objects
  46         // smaller than the size will be sent in a single request, while larger
  47         // objects will be split over multiple requests. The size will be rounded up
  48         // to the nearest multiple of 256K. If zero, chunking will be disabled and
  49         // the object will be uploaded in a single request.
  50         //
  51         // ChunkSize will default to a reasonable value. If you perform many concurrent
  52         // writes of small objects, you may wish set ChunkSize to a value that matches
  53         // your objects' sizes to avoid consuming large amounts of memory.
  54         //
  55         // ChunkSize must be set before the first Write call.
  56         ChunkSize int
  57
  58         // ProgressFunc can be used to monitor the progress of a large write.
  59         // operation. If ProgressFunc is not nil and writing requires multiple
  60         // calls to the underlying service (see
  61         // https://cloud.google.com/storage/docs/json_api/v1/how-tos/resumable-upload),
  62         // then ProgressFunc will be invoked after each call with the number of bytes of
  63         // content copied so far.
  64         //
  65         // ProgressFunc should return quickly without blocking.
  66         ProgressFunc func(int64)
  67
  68         ctx context.Context
  69         o   *ObjectHandle
  70
  71         opened bool
  72         pw     *io.PipeWriter
  73
  74         donec chan struct{} // closed after err and obj are set.
  75         obj   *ObjectAttrs
  76
  77         mu  sync.Mutex
  78         err error
  79 }
  80
  81 func (w *Writer) open() error {
  82         attrs := w.ObjectAttrs
  83         // Check the developer didn't change the object Name (this is unfortunate, but
  84         // we don't want to store an object under the wrong name).
  85         if attrs.Name != w.o.object {
  86                 return fmt.Errorf("storage: Writer.Name %q does not match object name %q", attrs.Name, w.o.object)
  87         }
  88         if !utf8.ValidString(attrs.Name) {
  89                 return fmt.Errorf("storage: object name %q is not valid UTF-8", attrs.Name)
  90         }
  91         if attrs.KMSKeyName != "" && w.o.encryptionKey != nil {
  92                 return errors.New("storage: cannot use KMSKeyName with a customer-supplied encryption key")
  93         }
  94         pr, pw := io.Pipe()
  95         w.pw = pw
  96         w.opened = true
  97
  98         go w.monitorCancel()
  99
 100         if w.ChunkSize < 0 {
 101                 return errors.New("storage: Writer.ChunkSize must be non-negative")
 102         }
 103         mediaOpts := []googleapi.MediaOption{
 104                 googleapi.ChunkSize(w.ChunkSize),
 105         }
 106         if c := attrs.ContentType; c != "" {
 107                 mediaOpts = append(mediaOpts, googleapi.ContentType(c))
 108         }
 109
 110         go func() {
 111                 defer close(w.donec)
 112
 113                 rawObj := attrs.toRawObject(w.o.bucket)
 114                 if w.SendCRC32C {
 115                         rawObj.Crc32c = encodeUint32(attrs.CRC32C)
 116                 }
 117                 if w.MD5 != nil {
 118                         rawObj.Md5Hash = base64.StdEncoding.EncodeToString(w.MD5)
 119                 }
 120                 call := w.o.c.raw.Objects.Insert(w.o.bucket, rawObj).
 121                         Media(pr, mediaOpts...).
 122                         Projection("full").
 123                         Context(w.ctx)
 124                 if w.ProgressFunc != nil {
 125                         call.ProgressUpdater(func(n, _ int64) { w.ProgressFunc(n) })
 126                 }
 127                 if attrs.KMSKeyName != "" {
 128                         call.KmsKeyName(attrs.KMSKeyName)
 129                 }
 130                 if attrs.PredefinedACL != "" {
 131                         call.PredefinedAcl(attrs.PredefinedACL)
 132                 }
 133                 if err := setEncryptionHeaders(call.Header(), w.o.encryptionKey, false); err != nil {
 134                         w.mu.Lock()
 135                         w.err = err
 136                         w.mu.Unlock()
 137                         pr.CloseWithError(err)
 138                         return
 139                 }
 140                 var resp *raw.Object
 141                 err := applyConds("NewWriter", w.o.gen, w.o.conds, call)
 142                 if err == nil {
 143                         if w.o.userProject != "" {
 144                                 call.UserProject(w.o.userProject)
 145                         }
 146                         setClientHeader(call.Header())
 147                         // If the chunk size is zero, then no chunking is done on the Reader,
 148                         // which means we cannot retry: the first call will read the data, and if
 149                         // it fails, there is no way to re-read.
 150                         if w.ChunkSize == 0 {
 151                                 resp, err = call.Do()
 152                         } else {
 153                                 // We will only retry here if the initial POST, which obtains a URI for
 154                                 // the resumable upload, fails with a retryable error. The upload itself
 155                                 // has its own retry logic.
 156                                 err = runWithRetry(w.ctx, func() error {
 157                                         var err2 error
 158                                         resp, err2 = call.Do()
 159                                         return err2
 160                                 })
 161                         }
 162                 }
 163                 if err != nil {
 164                         w.mu.Lock()
 165                         w.err = err
 166                         w.mu.Unlock()
 167                         pr.CloseWithError(err)
 168                         return
 169                 }
 170                 w.obj = newObject(resp)
 171         }()
 172         return nil
 173 }
 174
 175 // Write appends to w. It implements the io.Writer interface.
 176 //
 177 // Since writes happen asynchronously, Write may return a nil
 178 // error even though the write failed (or will fail). Always
 179 // use the error returned from Writer.Close to determine if
 180 // the upload was successful.
 181 func (w *Writer) Write(p []byte) (n int, err error) {
 182         w.mu.Lock()
 183         werr := w.err
 184         w.mu.Unlock()
 185         if werr != nil {
 186                 return 0, werr
 187         }
 188         if !w.opened {
 189                 if err := w.open(); err != nil {
 190                         return 0, err
 191                 }
 192         }
 193         n, err = w.pw.Write(p)
 194         if err != nil {
 195                 w.mu.Lock()
 196                 werr := w.err
 197                 w.mu.Unlock()
 198                 // Preserve existing functionality that when context is canceled, Write will return
 199                 // context.Canceled instead of "io: read/write on closed pipe". This hides the
 200                 // pipe implementation detail from users and makes Write seem as though it's an RPC.
 201                 if werr == context.Canceled || werr == context.DeadlineExceeded {
 202                         return n, werr
 203                 }
 204         }
 205         return n, err
 206 }
 207
 208 // Close completes the write operation and flushes any buffered data.
 209 // If Close doesn't return an error, metadata about the written object
 210 // can be retrieved by calling Attrs.
 211 func (w *Writer) Close() error {
 212         if !w.opened {
 213                 if err := w.open(); err != nil {
 214                         return err
 215                 }
 216         }
 217
 218         // Closing either the read or write causes the entire pipe to close.
 219         if err := w.pw.Close(); err != nil {
 220                 return err
 221         }
 222
 223         <-w.donec
 224         w.mu.Lock()
 225         defer w.mu.Unlock()
 226         return w.err
 227 }
 228
 229 // monitorCancel is intended to be used as a background goroutine. It monitors the
 230 // the context, and when it observes that the context has been canceled, it manually
 231 // closes things that do not take a context.
 232 func (w *Writer) monitorCancel() {
 233         select {
 234         case <-w.ctx.Done():
 235                 w.mu.Lock()
 236                 werr := w.ctx.Err()
 237                 w.err = werr
 238                 w.mu.Unlock()
 239
 240                 // Closing either the read or write causes the entire pipe to close.
 241                 w.CloseWithError(werr)
 242         case <-w.donec:
 243         }
 244 }
 245
 246 // CloseWithError aborts the write operation with the provided error.
 247 // CloseWithError always returns nil.
 248 //
 249 // Deprecated: cancel the context passed to NewWriter instead.
 250 func (w *Writer) CloseWithError(err error) error {
 251         if !w.opened {
 252                 return nil
 253         }
 254         return w.pw.CloseWithError(err)
 255 }
 256
 257 // Attrs returns metadata about a successfully-written object.
 258 // It's only valid to call it after Close returns nil.
 259 func (w *Writer) Attrs() *ObjectAttrs {
 260         return w.obj
 261 }