supporting heartbeat in cloudevents subscribe stream

Signed-off-by: Wei Liu <liuweixa@redhat.com>

Author: skeeey

pkg/cloudevents/generic/options/grpc/protocol/protocol.go

@@ -164,39 +164,7 @@ func (p *Protocol) OpenInbound(ctx context.Context) error {
 	}()
 
 	// start to monitor the server heartbeat
-	go func() {
-		// if serverHealthinessTimeout is disable, ignore the server heartbeat
-		if p.serverHealthinessTimeout == 0 {
-			for {
-				select {
-				case heartbeat := <-heartbeatCh:
-					logger.Debugf("heartbeat received %v", heartbeat)
-				case <-subCtx.Done():
-					// context canceled, exiting
-					return
-				}
-			}
-		}
-
-		// if no heartbeat was received duration the serverHealthinessTimeout, send the timeout error to
-		// reconnectErrorChan
-		for {
-			select {
-			case heartbeat := <-heartbeatCh:
-				logger.Debugf("heartbeat received %v", heartbeat)
-			case <-time.After(p.serverHealthinessTimeout):
-				select {
-				case p.reconnectErrorChan <- fmt.Errorf("stream timeout: no message received for %d seconds", p.serverHealthinessTimeout):
-				default:
-					// error chan is full, do nothing
-				}
-				return
-			case <-subCtx.Done():
-				// context canceled, exiting
-				return
-			}
-		}
-	}()
+	go startHeartbeatWatcher(ctx, heartbeatCh)
 
 	// Wait until external or internal context done
 	select {
@@ -229,3 +197,60 @@ func (p *Protocol) Close(ctx context.Context) error {
 	close(p.closeChan)
 	return nil
 }
+
+func startHeartbeatWatcher(
+	ctx context.Context,
+	heartbeatCh <-chan *pbv1.CloudEvent,
+) {
+	logger := cecontext.LoggerFrom(ctx)
+	// if serverHealthinessTimeout is disable, ignore the server heartbeat timeout
+	if p.serverHealthinessTimeout == 0 {
+		for {
+			select {
+			case msg, ok := <-heartbeatCh:
+				if !ok {
+					return
+				}
+				logger.Debugf("heartbeat received %v", msg)
+			case <-ctx.Done():
+				return
+			}
+		}
+	}
+
+	// if no heartbeat was received duration the serverHealthinessTimeout, send the timeout error to
+	// reconnectErrorChan
+	timeout := p.serverHealthinessTimeout
+	timer := time.NewTimer(timeout)
+	defer timer.Stop()
+
+	for {
+		select {
+		case msg, ok := <-heartbeatCh:
+			if !ok {
+				return
+			}
+			logger.Debugf("heartbeat received %v", msg)
+
+			// reset timer safely
+			if !timer.Stop() {
+				select {
+				case <-timer.C:
+				default:
+				}
+			}
+			timer.Reset(timeout)
+
+		case <-timer.C:
+			// timeout reached: no heartbeat received
+			select {
+			case p.reconnectErrorChan <- fmt.Errorf("stream timeout: no message received for %v", timeout):
+			default:
+			}
+			return
+
+		case <-ctx.Done():
+			return
+		}
+	}
+}