internal/streamer/streamer.go

@@ -64,6 +64,10 @@ const (
 	backpressureSlope     = 2.0 // Slope for linear drop probability increase
 )
 
+// messageTypeNotification is the type of a BGP NOTIFICATION message, which ends
+// the session with the peer.
+const messageTypeNotification = "NOTIFICATION"
+
 var (
 	errAlreadyRunning   = errors.New("streamer already running")
 	errUnexpectedStatus = errors.New("unexpected status code")
@@ -665,8 +669,8 @@ func (s *Streamer) readStream(
 }
 
 // processLine passes one line of the stream to the raw handler, parses it,
-// and dispatches an UPDATE or STATE message to the handlers. An error ends the
+// updates the peer set (updatePeers), and dispatches an UPDATE or STATE
-// connection.
+// message to the handlers. An error ends the connection.
 func (s *Streamer) processLine(ctx context.Context, line []byte) error {
 	// Update metrics with decompressed message size
 	s.updateMetrics(len(line))
@@ -711,6 +715,8 @@ func (s *Streamer) processLine(ctx context.Context, line []byte) error {
 	// Parse the timestamp
 	msg.ParsedTimestamp = time.Unix(int64(msg.Timestamp), 0).UTC()
 
+	peerCount := s.updatePeers(&msg)
+
 	// Process based on message type
 	switch msg.Type {
 	case ristypes.MessageTypeUpdate:
@@ -723,13 +729,6 @@ func (s *Streamer) processLine(ctx context.Context, line []byte) error {
 		// BGP keepalive messages - silently process
 		return nil
 	case "OPEN":
-		// BGP open messages - track peer as active
-		s.bgpPeersMu.Lock()
-		s.bgpPeers[msg.Peer] = true
-		peerCount := len(s.bgpPeers)
-		s.bgpPeersMu.Unlock()
-		s.metrics.SetBGPPeerCount(peerCount)
-
 		s.logger.Info("BGP session opened",
 			"peer", msg.Peer,
 			"peer_asn", msg.PeerASN,
@@ -737,14 +736,7 @@ func (s *Streamer) processLine(ctx context.Context, line []byte) error {
 		)
 
 		return nil
-	case "NOTIFICATION":
+	case messageTypeNotification:
-		// BGP notification messages (session closed)
-		s.bgpPeersMu.Lock()
-		delete(s.bgpPeers, msg.Peer)
-		peerCount := len(s.bgpPeers)
-		s.bgpPeersMu.Unlock()
-		s.metrics.SetBGPPeerCount(peerCount)
-
 		s.logger.Warn("BGP notification",
 			"peer", msg.Peer,
 			"peer_asn", msg.PeerASN,
@@ -766,6 +758,29 @@ func (s *Streamer) processLine(ctx context.Context, line []byte) error {
 	return s.dispatchMessage(ctx, &msg)
 }
 
+// updatePeers adds the peer msg came from to the peer set, or removes it if
+// msg ends the peer's session: a NOTIFICATION, or a STATE message whose state
+// is down. It records the size of the set as the BGP peer count and returns
+// it.
+func (s *Streamer) updatePeers(msg *ristypes.RISMessage) int {
+	sessionEnded := msg.Type == messageTypeNotification ||
+		(msg.Type == ristypes.MessageTypeState && msg.State == ristypes.StateDown)
+
+	s.bgpPeersMu.Lock()
+
+	if sessionEnded {
+		delete(s.bgpPeers, msg.Peer)
+	} else {
+		s.bgpPeers[msg.Peer] = true
+	}
+
+	peerCount := len(s.bgpPeers)
+	s.bgpPeersMu.Unlock()
+	s.metrics.SetBGPPeerCount(peerCount)
+
+	return peerCount
+}
+
 // dispatchMessage queues msg for each handler that wants it, unless
 // queueMessage drops it for that handler.
 func (s *Streamer) dispatchMessage(

internal/streamer/streamer_internal_test.go

@@ -518,6 +518,78 @@ func TestStreamReadsLongLine(t *testing.T) {
 	}
 }
 
+// TestPeerCountCountsEveryPeerHeardFrom has the server send an UPDATE from
+// stateDownPeer and a KEEPALIVE from another peer, with no OPEN from either,
+// and then a message that ends stateDownPeer's session: a NOTIFICATION, or a
+// STATE message whose state is down. The peer count must be 2 after the first
+// two messages and 1 after the third.
+func TestPeerCountCountsEveryPeerHeardFrom(t *testing.T) {
+	t.Parallel()
+
+	const (
+		otherPeer      = "192.0.2.2"
+		peerUpdateLine = `{"type":"ris_message","data":{"type":"UPDATE",` +
+			`"peer":"` + stateDownPeer + `",` +
+			`"peer_asn":"64496","timestamp":1700000000}}` + "\n"
+		otherPeerKeepaliveLine = `{"type":"ris_message","data":{"type":"KEEPALIVE",` +
+			`"peer":"` + otherPeer + `",` +
+			`"peer_asn":"64497","timestamp":1700000000}}` + "\n"
+		peerNotificationLine = `{"type":"ris_message","data":{"type":"NOTIFICATION",` +
+			`"peer":"` + stateDownPeer + `",` +
+			`"peer_asn":"64496","timestamp":1700000000}}` + "\n"
+		bothPeers = 2
+	)
+
+	sessionEndLines := map[string]string{
+		"NOTIFICATION": peerNotificationLine,
+		"STATE down":   stateDownLine,
+	}
+
+	for name, sessionEndLine := range sessionEndLines {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			tracker := metrics.New()
+			s := New(logger.New(), tracker)
+
+			streamLines(t, s, peerUpdateLine+otherPeerKeepaliveLine)
+
+			got := tracker.GetBGPPeerCount()
+			if got != bothPeers {
+				t.Fatalf("peer count is %d after messages from two peers, want %d",
+					got, bothPeers)
+			}
+
+			streamLines(t, s, sessionEndLine)
+
+			got = tracker.GetBGPPeerCount()
+			if got != 1 {
+				t.Errorf("peer count is %d after one of two peers' sessions "+
+					"ended, want 1", got)
+			}
+		})
+	}
+}
+
+// streamLines has a server send lines as its whole response and runs one
+// stream of s against it.
+func streamLines(t *testing.T, s *Streamer, lines string) {
+	t.Helper()
+
+	srv := httptest.NewServer(http.HandlerFunc(
+		func(w http.ResponseWriter, _ *http.Request) {
+			_, _ = io.WriteString(w, lines)
+		}))
+	defer srv.Close()
+
+	s.url = srv.URL
+
+	err := s.stream(context.Background())
+	if !errors.Is(err, errStreamEnded) {
+		t.Fatalf("stream returned %v, want %v", err, errStreamEnded)
+	}
+}
+
 // checkReconnect runs a streamer against a fake RIS Live server. The first
 // response is served by firstResponse; every later one carries one UPDATE
 // line and then stays open. The streamer must come back for a second