-
Notifications
You must be signed in to change notification settings - Fork 0
/
buffer.go
119 lines (107 loc) · 2.32 KB
/
buffer.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
package logwriter
import (
"bytes"
"io"
"os"
"time"
)
func NewBuffer(size int, interval time.Duration, w io.WriteCloser) io.WriteCloser {
return newBuffer(size, interval, w, time.Now)
}
func newBuffer(size int, interval time.Duration, w io.WriteCloser, now func() time.Time) io.WriteCloser {
return &Buffer{
Size: size,
Interval: interval,
Now: now,
w: w,
lastFlush: now(),
}
}
// Buffer implements buffered io.Writer like bufio.Writer, but specializing in log writer.
type Buffer struct {
Size int
Interval time.Duration
// Now() is a function returns current time like time.Time().
// This function used to inject time from outside.
Now func() time.Time
w io.WriteCloser
buf bytes.Buffer
err error
lastFlush time.Time
closed bool
}
func (b *Buffer) Write(p []byte) (n int, err error) {
n, err = b.write(p)
return
}
func (b *Buffer) Close() (err error) {
err = b.close()
return
}
func (b *Buffer) flush() {
if b.err != nil {
return
}
n := b.buf.Len()
var wrote int
if n > 0 {
wrote, b.err = b.w.Write(b.buf.Bytes())
}
b.lastFlush = b.Now()
b.buf.Reset()
if b.err == nil && wrote < n {
b.err = io.ErrShortWrite
}
return
}
func (b *Buffer) needFlush() bool {
expired := b.Now().Sub(b.lastFlush).Abs().Nanoseconds() >= b.Interval.Nanoseconds()
overflow := b.Size <= b.buf.Len()
return expired || overflow
}
func (b *Buffer) write(p []byte) (int, error) {
if b.closed {
return 0, os.ErrClosed
}
if b.Size <= len(p) {
return b.largeWrite(p)
}
return b.smallWrite(p)
}
func (b *Buffer) smallWrite(p []byte) (int, error) {
var n int
if b.err == nil && 0 < len(p) {
// This operation may require a buffer space of twice b.Size.
n, b.err = b.buf.Write(p)
}
if b.err == nil && b.needFlush() {
b.flush()
}
return n, b.err
}
func (b *Buffer) largeWrite(p []byte) (int, error) {
b.flush()
if b.err == nil && b.buf.Len() != 0 {
panic("bug: internal buffer is not empty even though it was flushed")
}
var wrote int
if b.err == nil {
wrote, b.err = b.w.Write(p)
b.lastFlush = b.Now()
}
if b.err == nil && wrote < len(p) {
b.err = io.ErrShortWrite
}
return wrote, b.err
}
func (b *Buffer) close() error {
if b.closed {
return os.ErrClosed
}
b.flush()
if b.err == nil {
b.err = b.w.Close()
b.closed = true
}
return b.err
}