@@ -29,10 +29,11 @@ func NewBatchWriterSize(w Writer, size int, intr time.Duration) *BatchWriter {
29
29
bufferSize := 10000
30
30
31
31
bw := & BatchWriter {
32
- w : w ,
33
- buf : make (chan Record , bufferSize ),
34
- signal : make (chan struct {}, bufferSize ),
35
- errors : make (chan error , 1 ),
32
+ w : w ,
33
+ buf : make (chan Record , bufferSize ),
34
+ signal : make (chan struct {}, bufferSize ),
35
+ flushSignal : make (chan struct {}, 0 ), // must be blocking write/read
36
+ errors : make (chan error , 1 ),
36
37
37
38
size : size ,
38
39
intr : intr ,
@@ -60,8 +61,9 @@ type BatchWriter struct {
60
61
size int
61
62
intr time.Duration
62
63
63
- ticker * time.Ticker
64
- signal chan struct {}
64
+ ticker * time.Ticker
65
+ signal chan struct {}
66
+ flushSignal chan struct {}
65
67
}
66
68
67
69
func (bw * BatchWriter ) writeLoop () {
@@ -77,11 +79,17 @@ func (bw *BatchWriter) writeLoop() {
77
79
}
78
80
case <- bw .ticker .C :
79
81
bw .flush ()
82
+ case _ , ok := <- bw .flushSignal :
83
+ if ok {
84
+ bw .flush ()
85
+ bw .flushSignal <- struct {}{}
86
+ }
80
87
}
81
88
}
82
89
}
83
90
84
- // flush writes everything in the current buffer
91
+ // flush writes everything in the current buffer. To prevent concurrent flushes,
92
+ // only write_loop() is allowed to call this.
85
93
func (bw * BatchWriter ) flush () {
86
94
numBuffered := len (bw .buf )
87
95
if numBuffered == 0 {
@@ -109,7 +117,8 @@ func (bw *BatchWriter) flush() {
109
117
// Flush forces all buffered records to be sent.
110
118
// If there is an error it will have been written to the Errors chan.
111
119
func (bw * BatchWriter ) Flush () {
112
- bw .flush ()
120
+ bw .flushSignal <- struct {}{} // ask to flush
121
+ <- bw .flushSignal // block until finish
113
122
}
114
123
115
124
// Close prevents future writes and flushes all currently buffered records.
@@ -118,6 +127,7 @@ func (bw *BatchWriter) Close() {
118
127
bw .ticker .Stop ()
119
128
close (bw .signal )
120
129
bw .Flush ()
130
+ close (bw .flushSignal )
121
131
close (bw .errors )
122
132
}
123
133
@@ -143,6 +153,6 @@ func (bw *BatchWriter) WriteRecords(rs ...Record) error {
143
153
for _ , r := range rs {
144
154
bw .buf <- r
145
155
}
146
- bw .signal <- struct {}{} // sigal that buffer changed
156
+ bw .signal <- struct {}{} // signal that buffer changed
147
157
return nil
148
158
}
0 commit comments