Project import generated by Copybara.

FolderOrigin-RevId: /usr/local/google/home/hines/copybara/temp/folder-destination3580118442091374722/.

Author: Googler

manager/manager.go

@@ -69,6 +69,9 @@ type Config struct {
 	Sync func(string)
 	// Timeout defines the optional duration to wait for a gRPC dial.
 	Timeout time.Duration
+	// ReceiveTimeout defines the optional duration to wait for receiving from
+	// a target. 0 means timeout is disabled.
+	ReceiveTimeout time.Duration
 	// Update will be invoked in response to gNMI updates.
 	Update func(string, *gpb.Notification)
 	// ConnectionManager is used to create gRPC connections.
@@ -99,6 +102,7 @@ type Manager struct {
 	sync              func(string)
 	testSync          func() // exposed for test synchronization
 	timeout           time.Duration
+	receiveTimeout    time.Duration
 	update            func(string, *gpb.Notification)
 
 	mu      sync.Mutex
@@ -134,6 +138,7 @@ func NewManager(cfg Config) (*Manager, error) {
 		targets:           make(map[string]*target),
 		testSync:          func() {},
 		timeout:           cfg.Timeout,
+		receiveTimeout:    cfg.ReceiveTimeout,
 		update:            cfg.Update,
 	}, nil
 }
@@ -191,11 +196,30 @@ func (m *Manager) createConn(ctx context.Context, name string, t *tpb.Target) (c
 	}
 }
 
-func (m *Manager) handleUpdates(name string, sc gpb.GNMI_SubscribeClient) error {
+func (m *Manager) handleUpdates(ctx context.Context, name string, sc gpb.GNMI_SubscribeClient) error {
 	defer m.testSync()
 	connected := false
+	var recvTimer *time.Timer
+	if m.receiveTimeout.Nanoseconds() > 0 {
+		recvTimer = time.NewTimer(m.receiveTimeout)
+		recvTimer.Stop()
+		go func() {
+			select {
+			case <-ctx.Done():
+			case <-recvTimer.C:
+				log.Errorf("Timed out waiting to receive from %q after %v", name, m.receiveTimeout)
+				m.Reconnect(name)
+			}
+		}()
+	}
 	for {
+		if recvTimer != nil {
+			recvTimer.Reset(m.receiveTimeout)
+		}
 		resp, err := sc.Recv()
+		if recvTimer != nil {
+			recvTimer.Stop()
+		}
 		if err != nil {
 			if m.reset != nil {
 				m.reset(name)
@@ -238,7 +262,7 @@ func (m *Manager) subscribe(ctx context.Context, name string, conn *grpc.ClientC
 	if err := sc.Send(cr); err != nil {
 		return fmt.Errorf("error sending subscription request to target %q: %v", name, err)
 	}
-	if err = m.handleUpdates(name, sc); err != nil {
+	if err = m.handleUpdates(ctx, name, sc); err != nil {
 		return fmt.Errorf("stream failed for target %q: %v", name, err)
 	}
 	return nil

manager/manager_test.go

@@ -554,6 +554,87 @@ func TestRetrySubscribe(t *testing.T) {
 	r.assertLast(t, "no recovery after remove", nil, nil, nil, 0)
 }
 
+func TestReceiveTimeout(t *testing.T) {
+	f := newFakeConnection(t)
+	addr, stop := newDevice(t, &fakeServer{})
+	defer stop()
+
+	origSubscribeClient := subscribeClient
+	defer func() { subscribeClient = origSubscribeClient }()
+	fc := newFakeSubscribeClient()
+	subscribeClient = func(ctx context.Context, conn *grpc.ClientConn) (gpb.GNMI_SubscribeClient, error) {
+		go func() {
+			select {
+			case <-ctx.Done():
+				fc.sendRecvErr() // Simulate stream closure.
+			}
+		}()
+		return fc, nil
+	}
+
+	r := record{}
+	m, err := NewManager(Config{
+		Credentials:    &fakeCreds{},
+		Timeout:        time.Minute,
+		ReceiveTimeout: 2 * time.Second,
+		Connect: func(s string) {
+			r.connects = append(r.connects, s)
+		},
+		ConnectionManager: f,
+		Sync: func(s string) {
+			r.syncs = append(r.syncs, s)
+		},
+		Reset: func(s string) {
+			r.resets = append(r.resets, s)
+		},
+		Update: func(_ string, g *gpb.Notification) {
+			r.updates = append(r.updates, g)
+		},
+	})
+	if err != nil {
+		t.Fatal("could not initialize Manager")
+	}
+	handleUpdateDoneOnce := make(chan struct{})
+	handleUpdateCalled := make(chan struct{}, 1)
+	m.testSync = func() {
+		select {
+		case <-handleUpdateDoneOnce:
+		default:
+			close(handleUpdateDoneOnce)
+		}
+		handleUpdateCalled <- struct{}{}
+	}
+
+	name := "device1"
+	err = m.Add(name, &tpb.Target{Addresses: []string{addr}}, validSubscribeRequest)
+	if err != nil {
+		t.Fatalf("got error adding: %v, want no error", err)
+	}
+	defer m.Remove(name)
+
+	_, ok := m.targets[name]
+	if !ok {
+		t.Fatalf("missing internal target")
+	}
+
+	<-handleUpdateDoneOnce // receive has timed out
+	<-handleUpdateCalled
+	r.assertLast(t, "after receive timeout", nil, nil, []string{name}, 0) // reset once
+
+	// Verify manager will try reconnecting to target
+	fc.sendSync()
+	<-handleUpdateCalled
+	r.assertLast(t, "sync after reconnect", []string{name}, []string{name}, nil, 0)
+
+	updateCount := 10
+	for i := 0; i < updateCount; i++ {
+		fc.sendUpdate()
+		<-handleUpdateCalled
+	}
+	r.assertLast(t, "updates sent after reconnect", nil, nil, nil, updateCount)
+	assertTargets(t, m, 1)
+}
+
 func TestRemoveDuringBackoff(t *testing.T) {
 	origBaseDelay, origMaxDelay := RetryBaseDelay, RetryMaxDelay
 	defer func() { RetryBaseDelay, RetryMaxDelay = origBaseDelay, origMaxDelay }()