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
|
// Package streamio contains wrappers intended for turning gRPC streams
// that send/receive messages with a []byte field into io.Writers and
// io.Readers.
//
package streamio
import (
"io"
"sync"
)
// NewReader turns receiver into an io.Reader. Errors from the receiver
// function are passed on unmodified. This means receiver should emit
// io.EOF when done.
func NewReader(receiver func() ([]byte, error)) io.Reader {
return &receiveReader{receiver: receiver}
}
type receiveReader struct {
receiver func() ([]byte, error)
data []byte
err error
}
func (rr *receiveReader) Read(p []byte) (int, error) {
if len(rr.data) == 0 && rr.err == nil {
rr.data, rr.err = rr.receiver()
}
n := copy(p, rr.data)
rr.data = rr.data[n:]
// We want to return any potential error only in case we have no
// buffered data left. Otherwise, it can happen that we do not relay
// bytes when the reader returns both data and an error.
if len(rr.data) == 0 {
return n, rr.err
}
return n, nil
}
// NewWriter turns sender into an io.Writer. The sender callback will
// receive []byte arguments of length at most WriteBufferSize.
func NewWriter(sender func(p []byte) error) io.Writer {
return &sendWriter{sender: sender}
}
// NewSyncWriter turns sender into an io.Writer. The sender callback will
// receive []byte arguments of length at most WriteBufferSize. All calls to the
// sender will be synchronized via the mutex.
func NewSyncWriter(m *sync.Mutex, sender func(p []byte) error) io.Writer {
return &sendWriter{
sender: func(p []byte) error {
m.Lock()
defer m.Unlock()
return sender(p)
},
}
}
// WriteBufferSize is the largest []byte that Write() will pass to its
// underlying send function.
var WriteBufferSize = 128 * 1024
type sendWriter struct {
sender func([]byte) error
}
func (sw *sendWriter) Write(p []byte) (int, error) {
var sent int
for len(p) > 0 {
chunkSize := len(p)
if chunkSize > WriteBufferSize {
chunkSize = WriteBufferSize
}
if err := sw.sender(p[:chunkSize]); err != nil {
return sent, err
}
sent += chunkSize
p = p[chunkSize:]
}
return sent, nil
}
|
