blob: 40eff4f018879bc42e3eb6113d36a57d8c284176 (
plain)
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
|
package limithandler_test
import (
"sync/atomic"
pb "gitlab.com/gitlab-org/gitaly/internal/middleware/limithandler/testpb"
"golang.org/x/net/context"
)
type server struct {
requestCount uint64
blockCh chan (struct{})
}
func (s *server) registerRequest() {
atomic.AddUint64(&s.requestCount, 1)
}
func (s *server) getRequestCount() int {
return int(atomic.LoadUint64(&s.requestCount))
}
func (s *server) Unary(ctx context.Context, in *pb.UnaryRequest) (*pb.UnaryResponse, error) {
s.registerRequest()
<-s.blockCh // Block to ensure concurrency
return &pb.UnaryResponse{Ok: true}, nil
}
func (s *server) StreamOutput(in *pb.StreamOutputRequest, stream pb.Test_StreamOutputServer) error {
s.registerRequest()
<-s.blockCh // Block to ensure concurrency
return stream.Send(&pb.StreamOutputResponse{Ok: true})
}
func (s *server) StreamInput(stream pb.Test_StreamInputServer) error {
// Read all the input
for {
if _, err := stream.Recv(); err != nil {
break
}
s.registerRequest()
}
<-s.blockCh // Block to ensure concurrency
return stream.SendAndClose(&pb.StreamInputResponse{Ok: true})
}
func (s *server) Bidirectional(stream pb.Test_BidirectionalServer) error {
// Read all the input
for {
if _, err := stream.Recv(); err != nil {
break
}
s.registerRequest()
}
<-s.blockCh // Block to ensure concurrency
return stream.Send(&pb.BidirectionalResponse{Ok: true})
}
|
