github.com/wlan0/docker@v1.5.0/pkg/broadcastwriter/broadcastwriter.go (about)

     1  package broadcastwriter
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"sync"
     7  	"time"
     8  
     9  	log "github.com/Sirupsen/logrus"
    10  	"github.com/docker/docker/pkg/jsonlog"
    11  )
    12  
    13  // BroadcastWriter accumulate multiple io.WriteCloser by stream.
    14  type BroadcastWriter struct {
    15  	sync.Mutex
    16  	buf      *bytes.Buffer
    17  	jsLogBuf *bytes.Buffer
    18  	streams  map[string](map[io.WriteCloser]struct{})
    19  }
    20  
    21  // AddWriter adds new io.WriteCloser for stream.
    22  // If stream is "", then all writes proceed as is. Otherwise every line from
    23  // input will be packed to serialized jsonlog.JSONLog.
    24  func (w *BroadcastWriter) AddWriter(writer io.WriteCloser, stream string) {
    25  	w.Lock()
    26  	if _, ok := w.streams[stream]; !ok {
    27  		w.streams[stream] = make(map[io.WriteCloser]struct{})
    28  	}
    29  	w.streams[stream][writer] = struct{}{}
    30  	w.Unlock()
    31  }
    32  
    33  // Write writes bytes to all writers. Failed writers will be evicted during
    34  // this call.
    35  func (w *BroadcastWriter) Write(p []byte) (n int, err error) {
    36  	created := time.Now().UTC()
    37  	w.Lock()
    38  	if writers, ok := w.streams[""]; ok {
    39  		for sw := range writers {
    40  			if n, err := sw.Write(p); err != nil || n != len(p) {
    41  				// On error, evict the writer
    42  				delete(writers, sw)
    43  			}
    44  		}
    45  	}
    46  	if w.jsLogBuf == nil {
    47  		w.jsLogBuf = new(bytes.Buffer)
    48  		w.jsLogBuf.Grow(1024)
    49  	}
    50  	w.buf.Write(p)
    51  	for {
    52  		line, err := w.buf.ReadString('\n')
    53  		if err != nil {
    54  			w.buf.Write([]byte(line))
    55  			break
    56  		}
    57  		for stream, writers := range w.streams {
    58  			if stream == "" {
    59  				continue
    60  			}
    61  			jsonLog := jsonlog.JSONLog{Log: line, Stream: stream, Created: created}
    62  			err = jsonLog.MarshalJSONBuf(w.jsLogBuf)
    63  			if err != nil {
    64  				log.Errorf("Error making JSON log line: %s", err)
    65  				continue
    66  			}
    67  			w.jsLogBuf.WriteByte('\n')
    68  			b := w.jsLogBuf.Bytes()
    69  			for sw := range writers {
    70  				if _, err := sw.Write(b); err != nil {
    71  					delete(writers, sw)
    72  				}
    73  			}
    74  		}
    75  		w.jsLogBuf.Reset()
    76  	}
    77  	w.jsLogBuf.Reset()
    78  	w.Unlock()
    79  	return len(p), nil
    80  }
    81  
    82  // Clean closes and removes all writers. Last non-eol-terminated part of data
    83  // will be saved.
    84  func (w *BroadcastWriter) Clean() error {
    85  	w.Lock()
    86  	for _, writers := range w.streams {
    87  		for w := range writers {
    88  			w.Close()
    89  		}
    90  	}
    91  	w.streams = make(map[string](map[io.WriteCloser]struct{}))
    92  	w.Unlock()
    93  	return nil
    94  }
    95  
    96  func New() *BroadcastWriter {
    97  	return &BroadcastWriter{
    98  		streams: make(map[string](map[io.WriteCloser]struct{})),
    99  		buf:     bytes.NewBuffer(nil),
   100  	}
   101  }