github.com/geofffranks/garden-linux@v0.0.0-20160715111146-26c893169cfa/process_tracker/writer/fan_in.go (about)

     1  package writer
     2  
     3  import (
     4  	"errors"
     5  	"io"
     6  	"sync"
     7  )
     8  
     9  type FanIn interface {
    10  	Write(data []byte) (int, error)
    11  	Close() error
    12  	AddSink(sink io.WriteCloser)
    13  	AddSource(source io.Reader)
    14  }
    15  
    16  func NewFanIn() FanIn {
    17  	return &fanIn{hasSink: make(chan struct{})}
    18  }
    19  
    20  type fanIn struct {
    21  	w      io.WriteCloser
    22  	closed bool
    23  	writeL sync.Mutex
    24  
    25  	hasSink chan struct{}
    26  }
    27  
    28  func (fw *fanIn) Write(data []byte) (int, error) {
    29  	<-fw.hasSink
    30  
    31  	fw.writeL.Lock()
    32  	defer fw.writeL.Unlock()
    33  
    34  	if fw.closed {
    35  		return 0, errors.New("write after close")
    36  	}
    37  
    38  	return fw.w.Write(data)
    39  }
    40  
    41  func (fw *fanIn) Close() error {
    42  	<-fw.hasSink
    43  
    44  	fw.writeL.Lock()
    45  	defer fw.writeL.Unlock()
    46  
    47  	if fw.closed {
    48  		return errors.New("already closed")
    49  	}
    50  
    51  	fw.closed = true
    52  
    53  	return fw.w.Close()
    54  }
    55  
    56  //AddSink can only be called once
    57  func (fw *fanIn) AddSink(sink io.WriteCloser) {
    58  	fw.w = sink
    59  
    60  	// Allow Write and Close to proceed.
    61  	close(fw.hasSink)
    62  }
    63  
    64  func (fw *fanIn) AddSource(source io.Reader) {
    65  	go func() {
    66  		_, err := io.Copy(fw, source)
    67  		if err == nil {
    68  			fw.Close()
    69  		}
    70  	}()
    71  }