// DoCancelAfterFirstResponse cancels the RPC after receiving the first message from the server.
func DoCancelAfterFirstResponse(tc testpb.TestServiceClient) {
ctx, cancel := context.WithCancel(context.Background())
stream, err := tc.FullDuplexCall(ctx)
if err != nil {
grpclog.Fatalf("%v.FullDuplexCall(_) = _, %v", tc, err)
}
respParam := []*testpb.ResponseParameters{
{
Size: proto.Int32(31415),
},
}
pl := clientNewPayload(testpb.PayloadType_COMPRESSABLE, 27182)
req := &testpb.StreamingOutputCallRequest{
ResponseType: testpb.PayloadType_COMPRESSABLE.Enum(),
ResponseParameters: respParam,
Payload: pl,
}
if err := stream.Send(req); err != nil {
grpclog.Fatalf("%v.Send(%v) = %v", stream, req, err)
}
if _, err := stream.Recv(); err != nil {
grpclog.Fatalf("%v.Recv() = %v", stream, err)
}
cancel()
if _, err := stream.Recv(); grpc.Code(err) != codes.Canceled {
grpclog.Fatalf("%v compleled with error code %d, want %d", stream, grpc.Code(err), codes.Canceled)
}
grpclog.Println("CancelAfterFirstResponse done")
}
func (t *http2Client) newStream(ctx context.Context, callHdr *CallHdr) *Stream {
fc := &inFlow{
limit: initialWindowSize,
conn: t.fc,
}
// TODO(zhaoq): Handle uint32 overflow of Stream.id.
s := &Stream{
id: t.nextID,
method: callHdr.Method,
sendCompress: callHdr.SendCompress,
buf: newRecvBuffer(),
fc: fc,
sendQuotaPool: newQuotaPool(int(t.streamSendQuota)),
headerChan: make(chan struct{}),
}
t.nextID += 2
s.windowHandler = func(n int) {
t.updateWindow(s, uint32(n))
}
// Make a stream be able to cancel the pending operations by itself.
s.ctx, s.cancel = context.WithCancel(ctx)
s.dec = &recvBufferReader{
ctx: s.ctx,
recv: s.buf,
}
return s
}
// DoCancelAfterBegin cancels the RPC after metadata has been sent but before payloads are sent.
func DoCancelAfterBegin(tc testpb.TestServiceClient) {
ctx, cancel := context.WithCancel(metadata.NewContext(context.Background(), testMetadata))
stream, err := tc.StreamingInputCall(ctx)
if err != nil {
grpclog.Fatalf("%v.StreamingInputCall(_) = _, %v", tc, err)
}
cancel()
_, err = stream.CloseAndRecv()
if grpc.Code(err) != codes.Canceled {
grpclog.Fatalf("%v.CloseAndRecv() got error code %d, want %d", stream, grpc.Code(err), codes.Canceled)
}
grpclog.Println("CancelAfterBegin done")
}
// ReadRows reads rows from a table. f is called for each row.
// If f returns false, the stream is shut down and ReadRows returns.
// f owns its argument, and f is called serially.
//
// By default, the yielded rows will contain all values in all cells.
// Use RowFilter to limit the cells returned.
func (t *Table) ReadRows(ctx context.Context, arg RowRange, f func(Row) bool, opts ...ReadOption) error {
req := &btspb.ReadRowsRequest{
TableName: t.c.fullTableName(t.table),
Target: &btspb.ReadRowsRequest_RowRange{arg.proto()},
}
for _, opt := range opts {
opt.set(req)
}
ctx, cancel := context.WithCancel(ctx) // for aborting the stream
defer cancel()
stream, err := t.c.client.ReadRows(ctx, req)
if err != nil {
return err
}
cr := new(chunkReader)
for {
res, err := stream.Recv()
if err == io.EOF {
break
}
if err != nil {
return err
}
if row := cr.process(res); row != nil {
if !f(row) {
// Cancel and drain stream.
cancel()
for {
if _, err := stream.Recv(); err != nil {
// The stream has ended. We don't return an error
// because the caller has intentionally interrupted the scan.
return nil
}
}
}
}
}
return nil
}
// operateHeader takes action on the decoded headers. It returns the current
// stream if there are remaining headers on the wire (in the following
// Continuation frame).
func (t *http2Server) operateHeaders(hDec *hpackDecoder, s *Stream, frame headerFrame, endStream bool, handle func(*Stream)) (pendingStream *Stream) {
defer func() {
if pendingStream == nil {
hDec.state = decodeState{}
}
}()
endHeaders, err := hDec.decodeServerHTTP2Headers(frame)
if s == nil {
// s has been closed.
return nil
}
if err != nil {
grpclog.Printf("transport: http2Server.operateHeader found %v", err)
if se, ok := err.(StreamError); ok {
t.controlBuf.put(&resetStream{s.id, statusCodeConvTab[se.Code]})
}
return nil
}
if endStream {
// s is just created by the caller. No lock needed.
s.state = streamReadDone
}
if !endHeaders {
return s
}
s.recvCompress = hDec.state.encoding
if hDec.state.timeoutSet {
s.ctx, s.cancel = context.WithTimeout(context.TODO(), hDec.state.timeout)
} else {
s.ctx, s.cancel = context.WithCancel(context.TODO())
}
pr := &peer.Peer{
Addr: t.conn.RemoteAddr(),
}
// Attach Auth info if there is any.
if t.authInfo != nil {
pr.AuthInfo = t.authInfo
}
s.ctx = peer.NewContext(s.ctx, pr)
// Cache the current stream to the context so that the server application
// can find out. Required when the server wants to send some metadata
// back to the client (unary call only).
s.ctx = newContextWithStream(s.ctx, s)
// Attach the received metadata to the context.
if len(hDec.state.mdata) > 0 {
s.ctx = metadata.NewContext(s.ctx, hDec.state.mdata)
}
s.dec = &recvBufferReader{
ctx: s.ctx,
recv: s.buf,
}
s.recvCompress = hDec.state.encoding
s.method = hDec.state.method
t.mu.Lock()
if t.state != reachable {
t.mu.Unlock()
return nil
}
if uint32(len(t.activeStreams)) >= t.maxStreams {
t.mu.Unlock()
t.controlBuf.put(&resetStream{s.id, http2.ErrCodeRefusedStream})
return nil
}
s.sendQuotaPool = newQuotaPool(int(t.streamSendQuota))
t.activeStreams[s.id] = s
t.mu.Unlock()
s.windowHandler = func(n int) {
t.updateWindow(s, uint32(n))
}
handle(s)
return nil
}
