github.com/karrick/gorill@v1.10.3/multiWriteCloserFanIn.go (about)

     1  package gorill
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  )
     7  
     8  const (
     9  	bufSize = 4096
    10  )
    11  
    12  // MultiWriteCloserFanIn is a structure that provides multiple io.WriteClosers to write to same underlying
    13  // io.WriteCloser.  When the final io.WriteCloser that MultiWriteCloserFanIn provides is closed, then the
    14  // underlying io.WriteCloser will be closed.
    15  type MultiWriteCloserFanIn struct {
    16  	iowc  io.WriteCloser
    17  	done  sync.WaitGroup
    18  	pLock *sync.Mutex
    19  	pDone *sync.WaitGroup
    20  }
    21  
    22  // NewMultiWriteCloserFanIn creates a MultiWriteCloserFanIn instance where writes to any of the provided
    23  // io.WriteCloser instances will be funneled to the underlying io.WriteCloser instance.  The client
    24  // ought to call Close on all provided io.WriteCloser instances, after which, MultiWriteCloserFanIn will
    25  // close the underlying io.WriteCloser.
    26  //
    27  //    func Example(largeBuf []byte) {
    28  //    	bb := NewNopCloseBufferSize(16384)
    29  //    	first := NewMultiWriteCloserFanIn(bb)
    30  //    	second := first.Add()
    31  //    	first.Write(largeBuf)
    32  //    	first.Close()
    33  //    	second.Write(largeBuf)
    34  //    	second.Close()
    35  //    }
    36  func NewMultiWriteCloserFanIn(iowc io.WriteCloser) *MultiWriteCloserFanIn {
    37  	var lock sync.Mutex
    38  	var done sync.WaitGroup
    39  	prime := &MultiWriteCloserFanIn{iowc: iowc, pLock: &lock, pDone: &done}
    40  	d := prime.Add()
    41  	go func() {
    42  		done.Wait()
    43  		iowc.Close()
    44  	}()
    45  	return d
    46  }
    47  
    48  // Add returns a new MultiWriteCloserFanIn that redirects all writes to the underlying
    49  // io.WriteCloser.  The client ought to call Close on the returned MultiWriteCloserFanIn to signify
    50  // intent to no longer Write to the MultiWriteCloserFanIn.
    51  func (fanin *MultiWriteCloserFanIn) Add() *MultiWriteCloserFanIn {
    52  	d := &MultiWriteCloserFanIn{iowc: fanin.iowc, pLock: fanin.pLock, pDone: fanin.pDone}
    53  	d.done.Add(1)
    54  	d.pDone.Add(1)
    55  	go func() {
    56  		d.done.Wait()
    57  		d.pDone.Done()
    58  	}()
    59  	return d
    60  }
    61  
    62  // Write copies the entire data slice to the underlying io.WriteCloser, ensuring no other
    63  // MultiWriteCloserFanIn can interrupt this one's writing.
    64  func (fanin *MultiWriteCloserFanIn) Write(data []byte) (int, error) {
    65  	fanin.pLock.Lock()
    66  	var err error
    67  	var written, m int
    68  	for err == nil && written < len(data) {
    69  		m, err = fanin.iowc.Write(data[written:])
    70  		written += m
    71  	}
    72  	fanin.pLock.Unlock()
    73  	return written, err
    74  }
    75  
    76  // Close marks the MultiWriteCloserFanIn as finished.  The last Close method invoked for a group of
    77  // MultiWriteCloserFanIn instances will trigger a close of the underlying io.WriteCloser.
    78  func (fanin *MultiWriteCloserFanIn) Close() error {
    79  	fanin.done.Done()
    80  	return nil
    81  }