github.com/squaremo/docker@v1.3.2-0.20150516120342-42cfc9554972/pkg/broadcastwriter/broadcastwriter.go (about)

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