lib/crunchrun/bufthenwrite.go

   1 // Copyright (C) The Arvados Authors. All rights reserved.
   2 //
   3 // SPDX-License-Identifier: AGPL-3.0
   4
   5 package crunchrun
   6
   7 import (
   8         "bytes"
   9         "io"
  10         "sync"
  11 )
  12
  13 type bufThenWrite struct {
  14         buf bytes.Buffer
  15         w   io.Writer
  16         mtx sync.Mutex
  17 }
  18
  19 func (btw *bufThenWrite) SetWriter(w io.Writer) error {
  20         btw.mtx.Lock()
  21         defer btw.mtx.Unlock()
  22         btw.w = w
  23         _, err := io.Copy(w, &btw.buf)
  24         return err
  25 }
  26
  27 func (btw *bufThenWrite) Write(p []byte) (int, error) {
  28         btw.mtx.Lock()
  29         defer btw.mtx.Unlock()
  30         if btw.w == nil {
  31                 btw.w = &btw.buf
  32         }
  33         return btw.w.Write(p)
  34 }