launchpad.net/~rogpeppe/juju-core/500-errgo-fix@v0.0.0-20140213181702-000000002356/utils/tailer/tailer.go (about)

     1  // Copyright 2013 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package tailer
     5  
     6  import (
     7  	"bufio"
     8  	"bytes"
     9  	"io"
    10  	"os"
    11  	"time"
    12  
    13  	"launchpad.net/errgo/errors"
    14  	"launchpad.net/tomb"
    15  )
    16  
    17  var mask = errors.Mask
    18  
    19  const (
    20  	bufferSize = 4096
    21  	polltime   = time.Second
    22  	delimiter  = '\n'
    23  )
    24  
    25  var (
    26  	delimiters = []byte{delimiter}
    27  )
    28  
    29  // TailerFilterFunc decides if a line shall be tailed (func is nil or
    30  // returns true) of shall be omitted (func returns false).
    31  type TailerFilterFunc func(line []byte) bool
    32  
    33  // Tailer reads an input line by line an tails them into the passed Writer.
    34  // The lines have to be terminated with a newline.
    35  type Tailer struct {
    36  	tomb        tomb.Tomb
    37  	readSeeker  io.ReadSeeker
    38  	reader      *bufio.Reader
    39  	writeCloser io.WriteCloser
    40  	writer      *bufio.Writer
    41  	lines       int
    42  	filter      TailerFilterFunc
    43  	bufferSize  int
    44  	polltime    time.Duration
    45  }
    46  
    47  // NewTailer starts a Tailer which reads strings from the passed
    48  // ReadSeeker line by line. If a filter function is specified the read
    49  // lines are filtered. The matching lines are written to the passed
    50  // Writer. The reading begins the specified number of matching lines
    51  // from the end.
    52  func NewTailer(readSeeker io.ReadSeeker, writer io.Writer, lines int, filter TailerFilterFunc) *Tailer {
    53  	return newTailer(readSeeker, writer, lines, filter, bufferSize, polltime)
    54  }
    55  
    56  // newTailer starts a Tailer like NewTailer but allows the setting of
    57  // the read buffer size and the time between pollings for testing.
    58  func newTailer(readSeeker io.ReadSeeker, writer io.Writer, lines int, filter TailerFilterFunc,
    59  	bufferSize int, polltime time.Duration) *Tailer {
    60  	t := &Tailer{
    61  		readSeeker: readSeeker,
    62  		reader:     bufio.NewReaderSize(readSeeker, bufferSize),
    63  		writer:     bufio.NewWriter(writer),
    64  		lines:      lines,
    65  		filter:     filter,
    66  		bufferSize: bufferSize,
    67  		polltime:   polltime,
    68  	}
    69  	go func() {
    70  		defer t.tomb.Done()
    71  		t.tomb.Kill(t.loop())
    72  	}()
    73  	return t
    74  }
    75  
    76  // Stop tells the tailer to stop working.
    77  func (t *Tailer) Stop() error {
    78  	t.tomb.Kill(nil)
    79  	return t.tomb.Wait()
    80  }
    81  
    82  // Wait waits until the tailer is stopped due to command
    83  // or an error. In case of an error it returns the reason.
    84  func (t *Tailer) Wait() error {
    85  	return t.tomb.Wait()
    86  }
    87  
    88  // Dead returns the channel that can be used to wait until
    89  // the tailer is stopped.
    90  func (t *Tailer) Dead() <-chan struct{} {
    91  	return t.tomb.Dead()
    92  }
    93  
    94  // Err returns a possible error.
    95  func (t *Tailer) Err() error {
    96  	return t.tomb.Err()
    97  }
    98  
    99  // loop writes the last lines based on the buffer size to the
   100  // writer and then polls for more data to write it to the
   101  // writer too.
   102  func (t *Tailer) loop() error {
   103  	// Position the readSeeker.
   104  	if err := t.seekLastLines(); err != nil {
   105  		return mask(err)
   106  	}
   107  
   108  	// Start polling.
   109  	// TODO(mue) 2013-12-06
   110  	// Handling of read-seeker/files being truncated during
   111  	// tailing is currently missing!
   112  	timer := time.NewTimer(0)
   113  	for {
   114  		select {
   115  		case <-t.tomb.Dying():
   116  			return nil
   117  		case <-timer.C:
   118  			for {
   119  				line, readErr := t.readLine()
   120  				_, writeErr := t.writer.Write(line)
   121  				if writeErr != nil {
   122  					return writeErr
   123  				}
   124  				if readErr != nil {
   125  					if readErr != io.EOF {
   126  						return readErr
   127  					}
   128  					break
   129  				}
   130  			}
   131  			if writeErr := t.writer.Flush(); writeErr != nil {
   132  				return writeErr
   133  			}
   134  			timer.Reset(t.polltime)
   135  		}
   136  	}
   137  }
   138  
   139  // seekLastLines sets the read position of the ReadSeeker to the
   140  // wanted number of filtered lines before the end.
   141  func (t *Tailer) seekLastLines() error {
   142  	offset, err := t.readSeeker.Seek(0, os.SEEK_END)
   143  	if err != nil {
   144  		return mask(err)
   145  	}
   146  	seekPos := int64(0)
   147  	found := 0
   148  	buffer := make([]byte, t.bufferSize)
   149  SeekLoop:
   150  	for offset > 0 {
   151  		// buffer contains the data left over from the
   152  		// previous iteration.
   153  		space := cap(buffer) - len(buffer)
   154  		if space < t.bufferSize {
   155  			// Grow buffer.
   156  			newBuffer := make([]byte, len(buffer), cap(buffer)*2)
   157  			copy(newBuffer, buffer)
   158  			buffer = newBuffer
   159  			space = cap(buffer) - len(buffer)
   160  		}
   161  		if int64(space) > offset {
   162  			// Use exactly the right amount of space if there's
   163  			// only a small amount remaining.
   164  			space = int(offset)
   165  		}
   166  		// Copy data remaining from last time to the end of the buffer,
   167  		// so we can read into the right place.
   168  		copy(buffer[space:cap(buffer)], buffer)
   169  		buffer = buffer[0 : len(buffer)+space]
   170  		offset -= int64(space)
   171  		_, err := t.readSeeker.Seek(offset, os.SEEK_SET)
   172  		if err != nil {
   173  			return mask(err)
   174  		}
   175  		_, err = io.ReadFull(t.readSeeker, buffer[0:space])
   176  		if err != nil {
   177  			return mask(err)
   178  		}
   179  
   180  		// Find the end of the last line in the buffer.
   181  		// This will discard any unterminated line at the end
   182  		// of the file.
   183  		end := bytes.LastIndex(buffer, delimiters)
   184  		if end == -1 {
   185  			// No end of line found - discard incomplete
   186  			// line and continue looking. If this happens
   187  			// at the beginning of the file, we don't care
   188  			// because we're going to stop anyway.
   189  			buffer = buffer[:0]
   190  			continue
   191  		}
   192  		end++
   193  		for {
   194  			start := bytes.LastIndex(buffer[0:end-1], delimiters)
   195  			if start == -1 && offset >= 0 {
   196  				break
   197  			}
   198  			start++
   199  			if t.isValid(buffer[start:end]) {
   200  				found++
   201  				if found >= t.lines {
   202  					seekPos = offset + int64(start)
   203  					break SeekLoop
   204  				}
   205  			}
   206  			end = start
   207  		}
   208  		// Leave the last line in buffer, as we don't know whether
   209  		// it's complete or not.
   210  		buffer = buffer[0:end]
   211  	}
   212  	// Final positioning.
   213  	t.readSeeker.Seek(seekPos, os.SEEK_SET)
   214  	return nil
   215  }
   216  
   217  // readLine reads the next valid line from the reader, even if it is
   218  // larger than the reader buffer.
   219  func (t *Tailer) readLine() ([]byte, error) {
   220  	for {
   221  		slice, err := t.reader.ReadSlice(delimiter)
   222  		if err == nil {
   223  			if t.isValid(slice) {
   224  				return slice, nil
   225  			}
   226  			continue
   227  		}
   228  		line := append([]byte(nil), slice...)
   229  		for err == bufio.ErrBufferFull {
   230  			slice, err = t.reader.ReadSlice(delimiter)
   231  			line = append(line, slice...)
   232  		}
   233  		switch err {
   234  		case nil:
   235  			if t.isValid(line) {
   236  				return line, nil
   237  			}
   238  		case io.EOF:
   239  			// EOF without delimiter, step back.
   240  			t.readSeeker.Seek(-int64(len(line)), os.SEEK_CUR)
   241  			return nil, err
   242  		default:
   243  			return nil, err
   244  		}
   245  	}
   246  }
   247  
   248  // isValid checks if the passed line is valid by checking if the
   249  // line has content, the filter function is nil or it returns true.
   250  func (t *Tailer) isValid(line []byte) bool {
   251  	if t.filter == nil {
   252  		return true
   253  	}
   254  	return t.filter(line)
   255  }