github.com/kunnos/engine@v1.13.1/daemon/logger/copier.go (about)

     1  package logger
     2  
     3  import (
     4  	"bytes"
     5  	"io"
     6  	"sync"
     7  	"time"
     8  
     9  	"github.com/sirupsen/logrus"
    10  )
    11  
    12  const (
    13  	bufSize  = 16 * 1024
    14  	readSize = 2 * 1024
    15  )
    16  
    17  // Copier can copy logs from specified sources to Logger and attach Timestamp.
    18  // Writes are concurrent, so you need implement some sync in your logger.
    19  type Copier struct {
    20  	// srcs is map of name -> reader pairs, for example "stdout", "stderr"
    21  	srcs      map[string]io.Reader
    22  	dst       Logger
    23  	copyJobs  sync.WaitGroup
    24  	closeOnce sync.Once
    25  	closed    chan struct{}
    26  }
    27  
    28  // NewCopier creates a new Copier
    29  func NewCopier(srcs map[string]io.Reader, dst Logger) *Copier {
    30  	return &Copier{
    31  		srcs:   srcs,
    32  		dst:    dst,
    33  		closed: make(chan struct{}),
    34  	}
    35  }
    36  
    37  // Run starts logs copying
    38  func (c *Copier) Run() {
    39  	for src, w := range c.srcs {
    40  		c.copyJobs.Add(1)
    41  		go c.copySrc(src, w)
    42  	}
    43  }
    44  
    45  func (c *Copier) copySrc(name string, src io.Reader) {
    46  	defer c.copyJobs.Done()
    47  	buf := make([]byte, bufSize)
    48  	n := 0
    49  	eof := false
    50  	msg := &Message{Source: name}
    51  
    52  	for {
    53  		select {
    54  		case <-c.closed:
    55  			return
    56  		default:
    57  			// Work out how much more data we are okay with reading this time.
    58  			upto := n + readSize
    59  			if upto > cap(buf) {
    60  				upto = cap(buf)
    61  			}
    62  			// Try to read that data.
    63  			if upto > n {
    64  				read, err := src.Read(buf[n:upto])
    65  				if err != nil {
    66  					if err != io.EOF {
    67  						logrus.Errorf("Error scanning log stream: %s", err)
    68  						return
    69  					}
    70  					eof = true
    71  				}
    72  				n += read
    73  			}
    74  			// If we have no data to log, and there's no more coming, we're done.
    75  			if n == 0 && eof {
    76  				return
    77  			}
    78  			// Break up the data that we've buffered up into lines, and log each in turn.
    79  			p := 0
    80  			for q := bytes.Index(buf[p:n], []byte{'\n'}); q >= 0; q = bytes.Index(buf[p:n], []byte{'\n'}) {
    81  				msg.Line = buf[p : p+q]
    82  				msg.Timestamp = time.Now().UTC()
    83  				msg.Partial = false
    84  				select {
    85  				case <-c.closed:
    86  					return
    87  				default:
    88  					if logErr := c.dst.Log(msg); logErr != nil {
    89  						logrus.Errorf("Failed to log msg %q for logger %s: %s", msg.Line, c.dst.Name(), logErr)
    90  					}
    91  				}
    92  				p += q + 1
    93  			}
    94  			// If there's no more coming, or the buffer is full but
    95  			// has no newlines, log whatever we haven't logged yet,
    96  			// noting that it's a partial log line.
    97  			if eof || (p == 0 && n == len(buf)) {
    98  				if p < n {
    99  					msg.Line = buf[p:n]
   100  					msg.Timestamp = time.Now().UTC()
   101  					msg.Partial = true
   102  					if logErr := c.dst.Log(msg); logErr != nil {
   103  						logrus.Errorf("Failed to log msg %q for logger %s: %s", msg.Line, c.dst.Name(), logErr)
   104  					}
   105  					p = 0
   106  					n = 0
   107  				}
   108  				if eof {
   109  					return
   110  				}
   111  			}
   112  			// Move any unlogged data to the front of the buffer in preparation for another read.
   113  			if p > 0 {
   114  				copy(buf[0:], buf[p:n])
   115  				n -= p
   116  			}
   117  		}
   118  	}
   119  }
   120  
   121  // Wait waits until all copying is done
   122  func (c *Copier) Wait() {
   123  	c.copyJobs.Wait()
   124  }
   125  
   126  // Close closes the copier
   127  func (c *Copier) Close() {
   128  	c.closeOnce.Do(func() {
   129  		close(c.closed)
   130  	})
   131  }