rpcz/streaming.go at master · Jille/rpcz · GitHub

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
120
121
122
123
124
125
126
127
128
129
130
package rpcz

import (
	"context"
	"io"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
	"google.golang.org/grpc/peer"
)

func StreamClientInterceptor(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) {
	if !sampleCall(method) {
		return streamer(ctx, desc, cc, method, opts...)
	}
	c := &capturedCall{}
	c.Start(ctx, metadataFromOutgoingContext(ctx), nil)
	c.Record(method)
	var p peer.Peer
	opts = append(opts, grpc.Peer(&p))
	stream, err := streamer(ctx, desc, cc, method, opts...)
	if err != nil {
		c.Complete(err, p.Addr, false, nil)
		return stream, err
	}
	if p, ok := peer.FromContext(stream.Context()); ok { // grpc.Peer() will only fill p after the entire call completes, so we have to fetch it from the context here.
		c.SetPeer(p.Addr)
	}
	return &clientStream{stream, c}, nil
}

type clientStream struct {
	parent grpc.ClientStream
	cc     *capturedCall
}

func (s *clientStream) Header() (metadata.MD, error) {
	return s.parent.Header()
}

func (s *clientStream) Trailer() metadata.MD {
	return s.parent.Trailer()
}

func (s *clientStream) Context() context.Context {
	return s.parent.Context()
}

func (s *clientStream) CloseSend() error {
	if err := s.parent.CloseSend(); err != nil {
		s.cc.Complete(err, nil, false, nil)
		return err
	}
	return nil
}

func (s *clientStream) SendMsg(m interface{}) error {
	s.cc.RecordMessage(m, false)
	return s.parent.SendMsg(m)
}

func (s *clientStream) RecvMsg(m interface{}) error {
	if err := s.parent.RecvMsg(m); err != nil {
		if err == io.EOF {
			s.cc.Complete(nil, nil, false, nil)
		} else {
			s.cc.Complete(err, nil, false, nil)
		}
		return err
	}
	s.cc.RecordMessage(m, true)
	return nil
}

func StreamServerInterceptor(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
	if !sampleCall(info.FullMethod) {
		return handler(srv, ss)
	}
	c := &capturedCall{}
	c.inbound = true
	if p, ok := peer.FromContext(ss.Context()); ok {
		c.peer = p.Addr
	}
	c.Start(ss.Context(), metadataFromIncomingContext(ss.Context()), nil)
	c.Record(info.FullMethod)
	err := handler(srv, &serverStream{ss, c})
	c.Complete(err, nil, false, nil)
	return err
}

type serverStream struct {
	parent grpc.ServerStream
	cc     *capturedCall
}

func (s *serverStream) SetHeader(md metadata.MD) error {
	return s.parent.SetHeader(md)
}

func (s *serverStream) SendHeader(md metadata.MD) error {
	return s.parent.SendHeader(md)
}

func (s *serverStream) SetTrailer(md metadata.MD) {
	s.parent.SetTrailer(md)
}

func (s *serverStream) Context() context.Context {
	return s.parent.Context()
}

func (s *serverStream) SendMsg(m interface{}) error {
	if err := s.parent.SendMsg(m); err != nil {
		s.cc.Complete(err, nil, true, m)
		return err
	}
	s.cc.RecordMessage(m, false)
	return nil
}

func (s *serverStream) RecvMsg(m interface{}) error {
	if err := s.parent.RecvMsg(m); err != nil {
		if err != io.EOF {
			s.cc.Complete(err, nil, false, nil)
		}
		return err
	}
	s.cc.RecordMessage(m, true)
	return nil
}