forked from DE-labtory/bifrost
-
Notifications
You must be signed in to change notification settings - Fork 0
/
stream.go
92 lines (73 loc) · 1.9 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
84
85
86
87
88
89
90
91
92
package bifrost
import (
"context"
"github.com/DE-labtory/bifrost/pb"
"google.golang.org/grpc"
)
type Stream interface {
Send(*pb.Envelope) error
Recv() (*pb.Envelope, error)
}
type StreamWrapper interface {
Stream
Close()
GetStream() Stream
}
type CStreamWrapper struct {
conn *grpc.ClientConn
client pb.StreamServiceClient
clientStream pb.StreamService_BifrostStreamClient
cancel context.CancelFunc
}
//client stream wrapper
func NewClientStreamWrapper(conn *grpc.ClientConn) (StreamWrapper, error) {
ctx, cf := context.WithCancel(context.Background())
streamServiceClient := pb.NewStreamServiceClient(conn)
clientStream, err := streamServiceClient.BifrostStream(ctx)
if err != nil {
return nil, err
}
return &CStreamWrapper{
cancel: cf,
conn: conn,
clientStream: clientStream,
client: streamServiceClient,
}, nil
}
func (csw *CStreamWrapper) GetStream() Stream {
return csw.clientStream
}
func (csw *CStreamWrapper) Send(envelope *pb.Envelope) error {
return csw.clientStream.Send(envelope)
}
func (csw *CStreamWrapper) Recv() (*pb.Envelope, error) {
return csw.clientStream.Recv()
}
func (csw *CStreamWrapper) Close() {
csw.conn.Close()
csw.clientStream.CloseSend()
csw.cancel()
}
//server stream wrapper
type SStreamWrapper struct {
serverStream pb.StreamService_BifrostStreamServer
cancel context.CancelFunc
}
func NewServerStreamWrapper(serverStream pb.StreamService_BifrostStreamServer, cancel context.CancelFunc) StreamWrapper {
return &SStreamWrapper{
cancel: cancel,
serverStream: serverStream,
}
}
func (ssw *SStreamWrapper) GetStream() Stream {
return ssw.serverStream
}
func (ssw *SStreamWrapper) Close() {
ssw.cancel()
}
func (ssw *SStreamWrapper) Send(envelope *pb.Envelope) error {
return ssw.serverStream.Send(envelope)
}
func (ssw *SStreamWrapper) Recv() (*pb.Envelope, error) {
return ssw.serverStream.Recv()
}