Merge pull request #1 from ipfs/mirror-writer-awesome

better implementation of mirrorwriter

option.go

@@ -7,7 +7,7 @@ import (
 )
 
 // Global writer group for logs to output to
-var WriterGroup = new(MirrorWriter)
+var WriterGroup = NewMirrorWriter()
 
 type Option func()
 

writer.go

@@ -1,66 +1,240 @@
 package log
 
 import (
+ "fmt"
  "io"
  "sync"
- "time"
 )
 
+var MaxWriterBuffer = 512 * 1024
+
+var log = Logger("eventlog")
+
 type MirrorWriter struct {
- writers []io.WriteCloser
- lk sync.Mutex
+ active bool
+ activelk sync.Mutex
+
+ // channel for incoming writers
+ writerAdd chan io.WriteCloser
+
+ // slices of writer/sync-channel pairs
+ writers []*bufWriter
+
+ // synchronization channel for incoming writes
+ msgSync chan []byte
+}
+
+type writerSync struct {
+ w io.WriteCloser
+ br chan []byte
+}
+
+func NewMirrorWriter() *MirrorWriter {
+ mw := &MirrorWriter{
+ msgSync: make(chan []byte, 64), // sufficiently large buffer to avoid callers waiting
+ writerAdd: make(chan io.WriteCloser),
+ }
+
+ go mw.logRoutine()
+
+ return mw
 }
 
 func (mw *MirrorWriter) Write(b []byte) (int, error) {
- mw.lk.Lock()
- // write to all writers, and nil out the broken ones.
- var dropped bool
- done := make(chan error, 1)
- for i, w := range mw.writers {
- go func(out chan error) {
- _, err := w.Write(b)
- out <- err
- }(done)
+ mycopy := make([]byte, len(b))
+ copy(mycopy, b)
+ mw.msgSync <- mycopy
+ return len(b), nil
+}
+
+func (mw *MirrorWriter) Close() error {
+ // it is up to the caller to ensure that write is not called during or
+ // after close is called.
+ close(mw.msgSync)
+ return nil
+}
+
+func (mw *MirrorWriter) doClose() {
+ for _, w := range mw.writers {
+ w.writer.Close()
+ }
+}
+
+func (mw *MirrorWriter) logRoutine() {
+ // rebind to avoid races on nilling out struct fields
+ msgSync := mw.msgSync
+ writerAdd := mw.writerAdd
+
+ defer mw.doClose()
+
+ for {
  select {
- case err := <-done:
- if err != nil {
- mw.writers[i].Close()
- mw.writers[i] = nil
- dropped = true
+ case b, ok := <-msgSync:
+ if !ok {
+ return
  }
- case <-time.After(time.Millisecond * 500):
- mw.writers[i].Close()
+
+ // write to all writers
+ dropped := mw.broadcastMessage(b)
+
+ // consolidate the slice
+ if dropped {
+ mw.clearDeadWriters()
+ }
+ case w := <-writerAdd:
+ mw.writers = append(mw.writers, newBufWriter(w))
+
+ mw.activelk.Lock()
+ mw.active = true
+ mw.activelk.Unlock()
+ }
+ }
+}
+
+// broadcastMessage sends the given message to every writer
+// if any writer is killed during the send, 'true' is returned
+func (mw *MirrorWriter) broadcastMessage(b []byte) bool {
+ var dropped bool
+ for i, w := range mw.writers {
+ _, err := w.Write(b)
+ if err != nil {
  mw.writers[i] = nil
  dropped = true
-
- // clear channel out
- done = make(chan error, 1)
  }
  }
+ return dropped
+}
 
- // consolidate the slice
- if dropped {
- writers := mw.writers
- mw.writers = nil
- for _, w := range writers {
- if w != nil {
- mw.writers = append(mw.writers, w)
- }
+func (mw *MirrorWriter) clearDeadWriters() {
+ writers := mw.writers
+ mw.writers = nil
+ for _, w := range writers {
+ if w != nil {
+ mw.writers = append(mw.writers, w)
  }
  }
- mw.lk.Unlock()
- return len(b), nil
+ if len(mw.writers) == 0 {
+ mw.activelk.Lock()
+ mw.active = false
+ mw.activelk.Unlock()
+ }
 }
 
 func (mw *MirrorWriter) AddWriter(w io.WriteCloser) {
- mw.lk.Lock()
- mw.writers = append(mw.writers, w)
- mw.lk.Unlock()
+ mw.writerAdd <- w
 }
 
 func (mw *MirrorWriter) Active() (active bool) {
- mw.lk.Lock()
- active = len(mw.writers) > 0
- mw.lk.Unlock()
+ mw.activelk.Lock()
+ active = mw.active
+ mw.activelk.Unlock()
  return
 }
+
+func newBufWriter(w io.WriteCloser) *bufWriter {
+ bw := &bufWriter{
+ writer: w,
+ incoming: make(chan []byte, 1),
+ }
+
+ go bw.loop()
+ return bw
+}
+
+type bufWriter struct {
+ writer io.WriteCloser
+
+ incoming chan []byte
+
+ deathLock sync.Mutex
+ dead bool
+}
+
+var errDeadWriter = fmt.Errorf("writer is dead")
+
+func (bw *bufWriter) Write(b []byte) (int, error) {
+ bw.deathLock.Lock()
+ dead := bw.dead
+ bw.deathLock.Unlock()
+ if dead {
+ if bw.incoming != nil {
+ close(bw.incoming)
+ bw.incoming = nil
+ }
+ return 0, errDeadWriter
+ }
+
+ bw.incoming <- b
+ return len(b), nil
+}
+
+func (bw *bufWriter) die() {
+ bw.deathLock.Lock()
+ bw.dead = true
+ bw.writer.Close()
+ bw.deathLock.Unlock()
+}
+
+func (bw *bufWriter) loop() {
+ bufsize := 0
+ bufBase := make([][]byte, 0, 16) // some initial memory
+ buffered := bufBase
+ nextCh := make(chan []byte)
+
+ var nextMsg []byte
+
+ go func() {
+ for b := range nextCh {
+ _, err := bw.writer.Write(b)
+ if err != nil {
+ log.Info("eventlog write error: %s", err)
+ bw.die()
+ return
+ }
+ }
+ }()
+
+ // collect and buffer messages
+ incoming := bw.incoming
+ for {
+ if nextMsg == nil || nextCh == nil {
+ // nextCh == nil implies we are 'dead' and draining the incoming channel
+ // until the caller notices and closes it for us
+ select {
+ case b, ok := <-incoming:
+ if !ok {
+ return
+ }
+ nextMsg = b
+ }
+ }
+
+ select {
+ case b, ok := <-incoming:
+ if !ok {
+ return
+ }
+ bufsize += len(b)
+ buffered = append(buffered, b)
+ if bufsize > MaxWriterBuffer {
+ // if we have too many messages buffered, kill the writer
+ bw.die()
+ close(nextCh)
+ nextCh = nil
+ // explicity keep going here to drain incoming
+ }
+ case nextCh <- nextMsg:
+ nextMsg = nil
+ if len(buffered) > 0 {
+ nextMsg = buffered[0]
+ buffered = buffered[1:]
+ bufsize -= len(nextMsg)
+ }
+
+ if len(buffered) == 0 {
+ // reset slice position
+ buffered = bufBase[:0]
+ }
+ }
+ }
+}