Send server side RST stream on deadline execeeded

Author: arjan-bal

internal/transport/client_stream.go

@@ -59,7 +59,7 @@ func (s *ClientStream) Read(n int) (mem.BufferSlice, error) {
 	return b, err
 }
 
-// Close closes the stream and popagates err to any readers.
+// Close closes the stream and propagates err to any readers.
 func (s *ClientStream) Close(err error) {
 	var (
 		rst     bool

internal/transport/http2_client.go

@@ -1238,7 +1238,7 @@ func (t *http2Client) handleRSTStream(f *http2.RSTStreamFrame) {
 	if statusCode == codes.Canceled {
 		if d, ok := s.ctx.Deadline(); ok && !d.After(time.Now()) {
 			// Our deadline was already exceeded, and that was likely the cause
-			// of this cancellation.  Alter the status code accordingly.
+			// of this cancellation. Alter the status code accordingly.
 			statusCode = codes.DeadlineExceeded
 		}
 	}

internal/transport/http2_server.go

@@ -598,6 +598,15 @@ func (t *http2Server) operateHeaders(ctx context.Context, frame *http2.MetaHeade
 	if len(t.activeStreams) == 1 {
 		t.idle = time.Time{}
 	}
+	// Start a timer to close the stream on reaching the deadline.
+	if timeoutSet {
+		timer := time.AfterFunc(timeout, func() {
+			t.closeStream(s, true, http2.ErrCodeCancel, false)
+		})
+		s.deadlineTimerCancel = func() {
+			timer.Stop()
+		}
+	}
 	t.mu.Unlock()
 	if channelz.IsOn() {
 		t.channelz.SocketMetrics.StreamsStarted.Add(1)
@@ -1268,6 +1277,9 @@ func (t *http2Server) Close(err error) {
 	channelz.RemoveEntry(t.channelz.ID)
 	// Cancel all active streams.
 	for _, s := range streams {
+		if s.deadlineTimerCancel != nil {
+			s.deadlineTimerCancel()
+		}
 		s.cancel()
 	}
 }
@@ -1306,6 +1318,10 @@ func (t *http2Server) finishStream(s *ServerStream, rst bool, rstCode http2.ErrC
 		return
 	}
 
+	if s.deadlineTimerCancel != nil {
+		s.deadlineTimerCancel()
+	}
+
 	hdr.cleanup = &cleanupStream{
 		streamID: s.id,
 		rst:      rst,
@@ -1324,7 +1340,13 @@ func (t *http2Server) closeStream(s *ServerStream, rst bool, rstCode http2.ErrCo
 	// called to interrupt the potential blocking on other goroutines.
 	s.cancel()
 
-	s.swapState(streamDone)
+	oldState := s.swapState(streamDone)
+	if oldState == streamDone {
+		return
+	}
+	if s.deadlineTimerCancel != nil {
+		s.deadlineTimerCancel()
+	}
 	t.deleteStream(s, eosReceived)
 
 	t.controlBuf.put(&cleanupStream{

internal/transport/server_stream.go

@@ -34,9 +34,10 @@ import (
 type ServerStream struct {
 	*Stream // Embed for common stream functionality.
 
-	st      internalServerTransport
-	ctxDone <-chan struct{}    // closed at the end of stream.  Cache of ctx.Done() (for performance)
-	cancel  context.CancelFunc // invoked at the end of stream to cancel ctx.
+	st                  internalServerTransport
+	ctxDone             <-chan struct{}    // closed at the end of stream.  Cache of ctx.Done() (for performance)
+	cancel              context.CancelFunc // invoked at the end of stream to cancel ctx.
+	deadlineTimerCancel func()             // Invoked at the end of stream.
 
 	// Holds compressor names passed in grpc-accept-encoding metadata from the
 	// client.

internal/transport/transport_test.go

@@ -2961,3 +2961,90 @@ func (s) TestReadMessageHeaderMultipleBuffers(t *testing.T) {
 		t.Errorf("bytesRead = %d, want = %d", bytesRead, headerLen)
 	}
 }
+
+// Tests a scenario when the client doesn't send an RST frame when the
+// configured deadline is reached. The test verifies that the server sends an
+// RST stream only after the deadline is reached.
+func (s) TestServerSendsRSTAfterDeadlineToMisbehavedClient(t *testing.T) {
+	server := setUpServerOnly(t, 0, &ServerConfig{}, suspended)
+	defer server.stop()
+	// Create a client that can override server stream quota.
+	mconn, err := net.Dial("tcp", server.lis.Addr().String())
+	if err != nil {
+		t.Fatalf("Clent failed to dial:%v", err)
+	}
+	defer mconn.Close()
+	if err := mconn.SetWriteDeadline(time.Now().Add(time.Second * 10)); err != nil {
+		t.Fatalf("Failed to set write deadline: %v", err)
+	}
+	if n, err := mconn.Write(clientPreface); err != nil || n != len(clientPreface) {
+		t.Fatalf("mconn.Write(clientPreface) = %d, %v, want %d, <nil>", n, err, len(clientPreface))
+	}
+	// rstTimeChan chan indicates that reader received a RSTStream from server.
+	rstTimeChan := make(chan time.Time, 1)
+	var mu sync.Mutex
+	framer := http2.NewFramer(mconn, mconn)
+	if err := framer.WriteSettings(); err != nil {
+		t.Fatalf("Error while writing settings: %v", err)
+	}
+	go func() { // Launch a reader for this misbehaving client.
+		for {
+			frame, err := framer.ReadFrame()
+			if err != nil {
+				return
+			}
+			switch frame := frame.(type) {
+			case *http2.PingFrame:
+				// Write ping ack back so that server's BDP estimation works right.
+				mu.Lock()
+				framer.WritePing(true, frame.Data)
+				mu.Unlock()
+			case *http2.RSTStreamFrame:
+				if frame.Header().StreamID != 1 || http2.ErrCode(frame.ErrCode) != http2.ErrCodeCancel {
+					t.Errorf("RST stream received with streamID: %d and code: %v, want streamID: 1 and code: http2.ErrCodeCancel", frame.Header().StreamID, http2.ErrCode(frame.ErrCode))
+				}
+				rstTimeChan <- time.Now()
+				return
+			default:
+				// Do nothing.
+			}
+		}
+	}()
+	// Create a stream.
+	var buf bytes.Buffer
+	henc := hpack.NewEncoder(&buf)
+	if err := henc.WriteField(hpack.HeaderField{Name: ":method", Value: "POST"}); err != nil {
+		t.Fatalf("Error while encoding header: %v", err)
+	}
+	if err := henc.WriteField(hpack.HeaderField{Name: ":path", Value: "foo"}); err != nil {
+		t.Fatalf("Error while encoding header: %v", err)
+	}
+	if err := henc.WriteField(hpack.HeaderField{Name: ":authority", Value: "localhost"}); err != nil {
+		t.Fatalf("Error while encoding header: %v", err)
+	}
+	if err := henc.WriteField(hpack.HeaderField{Name: "content-type", Value: "application/grpc"}); err != nil {
+		t.Fatalf("Error while encoding header: %v", err)
+	}
+	if err := henc.WriteField(hpack.HeaderField{Name: "grpc-timeout", Value: "10m"}); err != nil {
+		t.Fatalf("Error while encoding header: %v", err)
+	}
+	mu.Lock()
+	startTime := time.Now()
+	if err := framer.WriteHeaders(http2.HeadersFrameParam{StreamID: 1, BlockFragment: buf.Bytes(), EndHeaders: true}); err != nil {
+		mu.Unlock()
+		t.Fatalf("Error while writing headers: %v", err)
+	}
+	mu.Unlock()
+
+	// Test server behavior for deadline expiration.
+	var rstTime time.Time
+	select {
+	case <-time.After(5 * time.Second):
+		t.Fatalf("Test timed out.")
+	case rstTime = <-rstTimeChan:
+	}
+
+	if got, want := rstTime.Sub(startTime), 10*time.Millisecond; got < want {
+		t.Fatalf("RST frame received earlier than expected by duration: %v", want-got)
+	}
+}