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 }