/
stream.go
83 lines (71 loc) · 2.44 KB
/
stream.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
package grpcutil
import (
"bufio"
"bytes"
"io"
"github.com/gogo/protobuf/types"
)
// StreamingBytesServer represents a server for an rpc method of the form:
// rpc Foo(Bar) returns (stream google.protobuf.BytesValue) {}
type StreamingBytesServer interface {
Send(bytesValue *types.BytesValue) error
}
// StreamingBytesClient represents a client for an rpc method of the form:
// rpc Foo(Bar) returns (stream google.protobuf.BytesValue) {}
type StreamingBytesClient interface {
Recv() (*types.BytesValue, error)
}
// NewStreamingBytesReader returns an io.Reader for a StreamingBytesClient.
func NewStreamingBytesReader(streamingBytesClient StreamingBytesClient) io.Reader {
return &streamingBytesReader{streamingBytesClient: streamingBytesClient}
}
type streamingBytesReader struct {
streamingBytesClient StreamingBytesClient
buffer bytes.Buffer
}
func (s *streamingBytesReader) Read(p []byte) (int, error) {
// TODO this is doing an unneeded copy (unless go is smarter than I think it is)
if s.buffer.Len() == 0 {
value, err := s.streamingBytesClient.Recv()
if err != nil {
return 0, err
}
if _, err := s.buffer.Write(value.Value); err != nil {
return 0, err
}
}
return s.buffer.Read(p)
}
// NewStreamingBytesWriter returns an io.Writer for a StreamingBytesServer.
func NewStreamingBytesWriter(streamingBytesServer StreamingBytesServer) io.Writer {
return &streamingBytesWriter{streamingBytesServer}
}
type streamingBytesWriter struct {
streamingBytesServer StreamingBytesServer
}
func (s *streamingBytesWriter) Write(p []byte) (int, error) {
if len(p) == 0 {
return 0, nil
}
if err := s.streamingBytesServer.Send(&types.BytesValue{Value: p}); err != nil {
return 0, err
}
return len(p), nil
}
// WriteToStreamingBytesServer writes the data from the io.Reader to the StreamingBytesServer.
func WriteToStreamingBytesServer(reader io.Reader, streamingBytesServer StreamingBytesServer) error {
_, err := bufio.NewReader(reader).WriteTo(NewStreamingBytesWriter(streamingBytesServer))
return err
}
// WriteFromStreamingBytesClient writes from the StreamingBytesClient to the io.Writer.
func WriteFromStreamingBytesClient(streamingBytesClient StreamingBytesClient, writer io.Writer) error {
for bytesValue, err := streamingBytesClient.Recv(); err != io.EOF; bytesValue, err = streamingBytesClient.Recv() {
if err != nil {
return err
}
if _, err = writer.Write(bytesValue.Value); err != nil {
return err
}
}
return nil
}