Skip to content

Commit 254bccb

Browse files
easwarseaswars
authored
idle: decrement active call count for streaming RPCs only when the call completes (#6610)
1 parent b0a946c commit 254bccb

File tree

3 files changed

+112
-78
lines changed

3 files changed

+112
-78
lines changed

‎call.go

-5
Original file line numberDiff line numberDiff line change
@@ -27,11 +27,6 @@ import (
2727
//
2828
// All errors returned by Invoke are compatible with the status package.
2929
func (cc *ClientConn) Invoke(ctx context.Context, method string, args, reply any, opts ...CallOption) error {
30-
if err := cc.idlenessMgr.OnCallBegin(); err != nil {
31-
return err
32-
}
33-
defer cc.idlenessMgr.OnCallEnd()
34-
3530
// allow interceptor to see all applicable call options, which means those
3631
// configured as defaults from dial option as well as per-call options
3732
opts = combine(cc.dopts.callOptions, opts)

‎internal/idle/idle_e2e_test.go

+102-68
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import (
2222
"context"
2323
"errors"
2424
"fmt"
25+
"io"
2526
"strings"
2627
"testing"
2728
"time"
@@ -179,80 +180,113 @@ func (s) TestChannelIdleness_Enabled_NoActivity(t *testing.T) {
179180
// Tests the case where channel idleness is enabled by passing a small value for
180181
// idle_timeout. Verifies that a READY channel with an ongoing RPC stays READY.
181182
func (s) TestChannelIdleness_Enabled_OngoingCall(t *testing.T) {
182-
// Create a ClientConn with a short idle_timeout.
183-
r := manual.NewBuilderWithScheme("whatever")
184-
dopts := []grpc.DialOption{
185-
grpc.WithTransportCredentials(insecure.NewCredentials()),
186-
grpc.WithResolvers(r),
187-
grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
188-
grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
189-
}
190-
cc, err := grpc.Dial(r.Scheme()+":///test.server", dopts...)
191-
if err != nil {
192-
t.Fatalf("grpc.Dial() failed: %v", err)
193-
}
194-
t.Cleanup(func() { cc.Close() })
195-
196-
// Start a test backend which keeps a unary RPC call active by blocking on a
197-
// channel that is closed by the test later on. Also push an address update
198-
// via the resolver.
199-
blockCh := make(chan struct{})
200-
backend := &stubserver.StubServer{
201-
EmptyCallF: func(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
202-
<-blockCh
203-
return &testpb.Empty{}, nil
183+
tests := []struct {
184+
name string
185+
makeRPC func(ctx context.Context, client testgrpc.TestServiceClient) error
186+
}{
187+
{
188+
name: "unary",
189+
makeRPC: func(ctx context.Context, client testgrpc.TestServiceClient) error {
190+
if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
191+
return fmt.Errorf("EmptyCall RPC failed: %v", err)
192+
}
193+
return nil
194+
},
195+
},
196+
{
197+
name: "streaming",
198+
makeRPC: func(ctx context.Context, client testgrpc.TestServiceClient) error {
199+
stream, err := client.FullDuplexCall(ctx)
200+
if err != nil {
201+
t.Fatalf("FullDuplexCall RPC failed: %v", err)
202+
}
203+
if _, err := stream.Recv(); err != nil && err != io.EOF {
204+
t.Fatalf("stream.Recv() failed: %v", err)
205+
}
206+
return nil
207+
},
204208
},
205209
}
206-
if err := backend.StartServer(); err != nil {
207-
t.Fatalf("Failed to start backend: %v", err)
208-
}
209-
t.Cleanup(backend.Stop)
210-
r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
211-
212-
// Verify that the ClientConn moves to READY.
213-
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
214-
defer cancel()
215-
testutils.AwaitState(ctx, t, cc, connectivity.Ready)
216210

217-
// Spawn a goroutine which checks expected state transitions and idleness
218-
// channelz trace events. It eventually closes `blockCh`, thereby unblocking
219-
// the server RPC handler and the unary call below.
220-
errCh := make(chan error, 1)
221-
go func() {
222-
defer close(blockCh)
223-
// Verify that the ClientConn stays in READY.
224-
sCtx, sCancel := context.WithTimeout(ctx, 3*defaultTestShortIdleTimeout)
225-
defer sCancel()
226-
testutils.AwaitNoStateChange(sCtx, t, cc, connectivity.Ready)
227-
228-
// Verify that there are no idleness related channelz events.
229-
if err := channelzTraceEventNotFound(ctx, "entering idle mode"); err != nil {
230-
errCh <- err
231-
return
232-
}
233-
if err := channelzTraceEventNotFound(ctx, "exiting idle mode"); err != nil {
234-
errCh <- err
235-
return
236-
}
211+
for _, test := range tests {
212+
t.Run(test.name, func(t *testing.T) {
213+
// Create a ClientConn with a short idle_timeout.
214+
r := manual.NewBuilderWithScheme("whatever")
215+
dopts := []grpc.DialOption{
216+
grpc.WithTransportCredentials(insecure.NewCredentials()),
217+
grpc.WithResolvers(r),
218+
grpc.WithIdleTimeout(defaultTestShortIdleTimeout),
219+
grpc.WithDefaultServiceConfig(`{"loadBalancingConfig": [{"round_robin":{}}]}`),
220+
}
221+
cc, err := grpc.Dial(r.Scheme()+":///test.server", dopts...)
222+
if err != nil {
223+
t.Fatalf("grpc.Dial() failed: %v", err)
224+
}
225+
t.Cleanup(func() { cc.Close() })
226+
227+
// Start a test backend which keeps a unary RPC call active by blocking on a
228+
// channel that is closed by the test later on. Also push an address update
229+
// via the resolver.
230+
blockCh := make(chan struct{})
231+
backend := &stubserver.StubServer{
232+
EmptyCallF: func(ctx context.Context, in *testpb.Empty) (*testpb.Empty, error) {
233+
<-blockCh
234+
return &testpb.Empty{}, nil
235+
},
236+
FullDuplexCallF: func(stream testgrpc.TestService_FullDuplexCallServer) error {
237+
<-blockCh
238+
return nil
239+
},
240+
}
241+
if err := backend.StartServer(); err != nil {
242+
t.Fatalf("Failed to start backend: %v", err)
243+
}
244+
t.Cleanup(backend.Stop)
245+
r.UpdateState(resolver.State{Addresses: []resolver.Address{{Addr: backend.Address}}})
246+
247+
// Verify that the ClientConn moves to READY.
248+
ctx, cancel := context.WithTimeout(context.Background(), defaultTestTimeout)
249+
defer cancel()
250+
testutils.AwaitState(ctx, t, cc, connectivity.Ready)
251+
252+
// Spawn a goroutine which checks expected state transitions and idleness
253+
// channelz trace events.
254+
errCh := make(chan error, 1)
255+
go func() {
256+
defer close(blockCh)
257+
258+
// Verify that the ClientConn stays in READY.
259+
sCtx, sCancel := context.WithTimeout(ctx, 3*defaultTestShortIdleTimeout)
260+
defer sCancel()
261+
if cc.WaitForStateChange(sCtx, connectivity.Ready) {
262+
errCh <- fmt.Errorf("state changed from %q to %q when no state change was expected", connectivity.Ready, cc.GetState())
263+
return
264+
}
237265

238-
// Unblock the unary RPC on the server.
239-
errCh <- nil
240-
}()
266+
// Verify that there are no idleness related channelz events.
267+
//
268+
// TODO: Improve the checks here. If these log strings are
269+
// changed in the code, these checks will continue to pass.
270+
if err := channelzTraceEventNotFound(ctx, "entering idle mode"); err != nil {
271+
errCh <- err
272+
return
273+
}
274+
errCh <- channelzTraceEventNotFound(ctx, "exiting idle mode")
275+
}()
241276

242-
// Make a unary RPC that blocks on the server, thereby ensuring that the
243-
// count of active RPCs on the client is non-zero.
244-
client := testgrpc.NewTestServiceClient(cc)
245-
if _, err := client.EmptyCall(ctx, &testpb.Empty{}); err != nil {
246-
t.Errorf("EmptyCall RPC failed: %v", err)
247-
}
277+
if err := test.makeRPC(ctx, testgrpc.NewTestServiceClient(cc)); err != nil {
278+
t.Fatalf("%s rpc failed: %v", test.name, err)
279+
}
248280

249-
select {
250-
case err := <-errCh:
251-
if err != nil {
252-
t.Fatal(err)
253-
}
254-
case <-ctx.Done():
255-
t.Fatalf("Timeout when trying to verify that an active RPC keeps channel from moving to IDLE")
281+
select {
282+
case err := <-errCh:
283+
if err != nil {
284+
t.Fatal(err)
285+
}
286+
case <-ctx.Done():
287+
t.Fatalf("Timeout when trying to verify that an active RPC keeps channel from moving to IDLE")
288+
}
289+
})
256290
}
257291
}
258292

‎stream.go

+10-5
Original file line numberDiff line numberDiff line change
@@ -158,11 +158,6 @@ type ClientStream interface {
158158
// If none of the above happen, a goroutine and a context will be leaked, and grpc
159159
// will not call the optionally-configured stats handler with a stats.End message.
160160
func (cc *ClientConn) NewStream(ctx context.Context, desc *StreamDesc, method string, opts ...CallOption) (ClientStream, error) {
161-
if err := cc.idlenessMgr.OnCallBegin(); err != nil {
162-
return nil, err
163-
}
164-
defer cc.idlenessMgr.OnCallEnd()
165-
166161
// allow interceptor to see all applicable call options, which means those
167162
// configured as defaults from dial option as well as per-call options
168163
opts = combine(cc.dopts.callOptions, opts)
@@ -179,6 +174,16 @@ func NewClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, meth
179174
}
180175

181176
func newClientStream(ctx context.Context, desc *StreamDesc, cc *ClientConn, method string, opts ...CallOption) (_ ClientStream, err error) {
177+
// Start tracking the RPC for idleness purposes. This is where a stream is
178+
// created for both streaming and unary RPCs, and hence is a good place to
179+
// track active RPC count.
180+
if err := cc.idlenessMgr.OnCallBegin(); err != nil {
181+
return nil, err
182+
}
183+
// Add a calloption, to decrement the active call count, that gets executed
184+
// when the RPC completes.
185+
opts = append([]CallOption{OnFinish(func(error) { cc.idlenessMgr.OnCallEnd() })}, opts...)
186+
182187
if md, added, ok := metadata.FromOutgoingContextRaw(ctx); ok {
183188
// validate md
184189
if err := imetadata.Validate(md); err != nil {

0 commit comments

Comments
 (0)