Welcome to mirror list, hosted at ThFree Co, Russian Federation.

testhelper_test.go « limithandler « middleware « grpc « internal - gitlab.com/gitlab-org/gitaly.git - Unnamed repository; edit this file 'description' to name the repository.
summaryrefslogtreecommitdiff
blob: 137c47cce2b1a3de88a6055c063856cdc2dd207d (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
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
package limithandler_test

import (
	"context"
	"io"
	"sync/atomic"

	"google.golang.org/grpc/interop/grpc_testing"
)

type server struct {
	grpc_testing.UnimplementedTestServiceServer
	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) UnaryCall(
	ctx context.Context,
	in *grpc_testing.SimpleRequest,
) (*grpc_testing.SimpleResponse, error) {
	s.registerRequest()

	<-s.blockCh // Block to ensure concurrency

	return &grpc_testing.SimpleResponse{
		Payload: &grpc_testing.Payload{
			Body: []byte("success"),
		},
	}, nil
}

func (s *server) StreamingOutputCall(
	in *grpc_testing.StreamingOutputCallRequest,
	stream grpc_testing.TestService_StreamingOutputCallServer,
) error {
	s.registerRequest()

	<-s.blockCh // Block to ensure concurrency

	return stream.Send(&grpc_testing.StreamingOutputCallResponse{
		Payload: &grpc_testing.Payload{
			Body: []byte("success"),
		},
	})
}

func (s *server) StreamingInputCall(stream grpc_testing.TestService_StreamingInputCallServer) error {
	// Read all the input
	for {
		if _, err := stream.Recv(); err != nil {
			if err != io.EOF {
				return err
			}
			break
		}

		s.registerRequest()
	}

	<-s.blockCh // Block to ensure concurrency

	return stream.SendAndClose(&grpc_testing.StreamingInputCallResponse{
		AggregatedPayloadSize: 9000,
	})
}

func (s *server) FullDuplexCall(stream grpc_testing.TestService_FullDuplexCallServer) error {
	// Read all the input
	for {
		if _, err := stream.Recv(); err != nil {
			if err != io.EOF {
				return err
			}
			break
		}

		s.registerRequest()
	}

	<-s.blockCh // Block to ensure concurrency

	return stream.Send(&grpc_testing.StreamingOutputCallResponse{
		Payload: &grpc_testing.Payload{
			Body: []byte("success"),
		},
	})
}