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