github.com/haraldrudell/parl@v0.4.176/pio/tee-writer.go (about)

     1  /*
     2  © 2021–present Harald Rudell <harald.rudell@gmail.com> (https://haraldrudell.github.io/haraldrudell/)
     3  ISC License
     4  */
     5  
     6  // TeeWriter is a writer that copies its writes to one or more other writers.
     7  package pio
     8  
     9  import (
    10  	"io"
    11  	"sync"
    12  	"sync/atomic"
    13  
    14  	"github.com/haraldrudell/parl"
    15  	"github.com/haraldrudell/parl/perrors"
    16  )
    17  
    18  // TeeWriter is a writer that copies its writes to one or more other writers.
    19  type TeeWriter struct {
    20  	closeCallback func() (err error)
    21  	writers       []io.Writer
    22  	isClosed      atomic.Bool
    23  	wg            sync.WaitGroup
    24  }
    25  
    26  // TeeWriter is a writer that copies its writes to one or more other writers.
    27  func NewTeeWriter(closeCallback func() (err error), writers ...io.Writer) (teeWriter io.WriteCloser) {
    28  	length := len(writers)
    29  	if length == 0 {
    30  		panic(perrors.NewPF("Must have one or more writers, writers is empty"))
    31  	}
    32  	t := TeeWriter{closeCallback: closeCallback, writers: make([]io.Writer, length)}
    33  	for i, w := range writers {
    34  		if w == nil {
    35  			panic(perrors.ErrorfPF("Writers#%d nil", i))
    36  		}
    37  		t.writers[i] = w
    38  	}
    39  	t.wg.Add(1)
    40  	return &t
    41  }
    42  
    43  func (tw *TeeWriter) Write(p []byte) (n int, err error) {
    44  	if tw.isClosed.Load() {
    45  		err = perrors.NewPF("Write after Close")
    46  		return
    47  	}
    48  	length := len(p)
    49  	for _, writer := range tw.writers {
    50  		written := 0
    51  		for written < length {
    52  			n, err = writer.Write(p)
    53  			written += n
    54  			if err != nil {
    55  				return // write error return
    56  			}
    57  		}
    58  	}
    59  	return // good write return
    60  }
    61  
    62  func (w *TeeWriter) Close() (err error) {
    63  
    64  	// prevent multiple Close invocations
    65  	if !w.isClosed.CompareAndSwap(false, true) {
    66  		err = perrors.NewPF("Second Close invocation")
    67  		return
    68  	}
    69  
    70  	// invoke callback if there is one
    71  	if w.closeCallback != nil {
    72  		err = w.invokeCallback()
    73  	}
    74  
    75  	return
    76  }
    77  
    78  func (w *TeeWriter) invokeCallback() (err error) {
    79  	defer parl.RecoverErr(func() parl.DA { return parl.A() }, &err)
    80  
    81  	err = w.closeCallback()
    82  
    83  	return
    84  }