github.com/u-root/u-root@v7.0.1-0.20200915234505-ad7babab0a8e+incompatible/pkg/ts/ts.go (about)

     1  // Copyright 2020 the u-root Authors. All rights reserved
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Package ts contains a Transform to prepend a timestamp in front of each line.
     6  package ts
     7  
     8  import (
     9  	"bytes"
    10  	"fmt"
    11  	"io"
    12  	"time"
    13  )
    14  
    15  // PrependTimestamp is an io.Reader which prepends a timestamp on each line.
    16  type PrependTimestamp struct {
    17  	R         io.Reader
    18  	StartTime time.Time
    19  	Format    func(startTime time.Time) string
    20  
    21  	// ResetTimeOnNextRead is set to force the StartTime to reset to the
    22  	// current time when data is next read. It is useful to set this to
    23  	// true initially for all times to be relative to the first line.
    24  	ResetTimeOnNextRead bool
    25  
    26  	// noPrintTS is true to indicate the timestamp should be printed
    27  	// in front of the next byte.
    28  	noPrintTS bool
    29  
    30  	// buf is a Buffer to store data in case the buffer passed to
    31  	// PrependTimestamp.Read is not large enough.
    32  	buf []byte
    33  
    34  	// isEOF indicates the wrapped reader R returned an io.EOF. All future
    35  	// calls to PrependTimestamp.Read should just read from buf until it is
    36  	// empty, then return io.EOF.
    37  	isEOF bool
    38  }
    39  
    40  // New creates a PrependTimestamp with default settings.
    41  func New(r io.Reader) *PrependTimestamp {
    42  	return &PrependTimestamp{
    43  		R:         r,
    44  		StartTime: time.Now(),
    45  		Format:    DefaultFormat,
    46  	}
    47  }
    48  
    49  // Read prepends a timestamp on each line.
    50  func (t *PrependTimestamp) Read(p []byte) (n int, err error) {
    51  	// Empty the buffer first.
    52  	n = copy(p, t.buf)
    53  	t.buf = t.buf[n:]
    54  	if len(t.buf) != 0 {
    55  		return // Buffer not yet empty.
    56  	}
    57  	if t.isEOF {
    58  		err = io.EOF
    59  		return // Buffer empty and reached EOF.
    60  	}
    61  
    62  	// Read new data.
    63  	var m int
    64  	scratch := p[n:]
    65  	m, err = t.R.Read(scratch)
    66  	scratch = scratch[:m]
    67  	if m == 0 {
    68  		return
    69  	}
    70  	if t.ResetTimeOnNextRead {
    71  		t.StartTime = time.Now()
    72  		t.ResetTimeOnNextRead = false
    73  	}
    74  	if err == io.EOF {
    75  		// Do not return EOF immediately because there may be more than
    76  		// one calls to PrependTimestamp.Read.
    77  		t.isEOF = true
    78  		err = nil
    79  	}
    80  
    81  	// Generate the timestamp.
    82  	lfts := []byte("\n" + t.Format(t.StartTime))
    83  	ts := lfts[1:]
    84  
    85  	// Insert timestamps after newlines.
    86  	t.buf = bytes.ReplaceAll(scratch, []byte{'\n'}, lfts)
    87  
    88  	// If the input ends in a newline, defer the insertion of the timestamp
    89  	// until the next byte is read.
    90  	if !t.noPrintTS {
    91  		t.buf = append(ts, t.buf...)
    92  		t.noPrintTS = true
    93  	}
    94  	if scratch[len(scratch)-1] == '\n' {
    95  		t.buf = t.buf[:len(t.buf)-len(ts)]
    96  		t.noPrintTS = false
    97  	}
    98  
    99  	// Empty new data from the buffer.
   100  	m = copy(p[n:], t.buf)
   101  	n += m
   102  	t.buf = t.buf[m:]
   103  	if len(t.buf) == 0 && t.isEOF {
   104  		err = io.EOF // Buffer empty and reached EOF.
   105  	}
   106  	return
   107  }
   108  
   109  // DefaultFormat formats in seconds since the startTime. Ex: [12.3456s]
   110  func DefaultFormat(startTime time.Time) string {
   111  	return fmt.Sprintf("[%06.4fs] ", time.Since(startTime).Seconds())
   112  }
   113  
   114  // NewRelativeFormat returns a format function which formats in seconds since
   115  // the previous line. Ex: [+1.0050s]
   116  func NewRelativeFormat() func(time.Time) string {
   117  	firstLine := true
   118  	var lastTime time.Time
   119  	return func(startTime time.Time) string {
   120  		if firstLine {
   121  			firstLine = false
   122  			lastTime = startTime
   123  		}
   124  		curTime := time.Now()
   125  		s := fmt.Sprintf("[+%06.4fs] ", curTime.Sub(lastTime).Seconds())
   126  		lastTime = curTime
   127  		return s
   128  	}
   129  }