From 5c979fea68b6000ec5339fe85dfbf21bc3977921 Mon Sep 17 00:00:00 2001
From: Alistair King <alistair@kentik.com>
Date: Tue, 3 Nov 2020 11:44:13 -0500
Subject: [PATCH] Use MonitorTableRequest.Name to filter by peer

Filter returned Paths based on provided peer address (similar to
GetTable). This should improve performance when using MonitorTable
with ADJ_IN and Current for a single peer.
---
 pkg/server/bmp.go     |  4 ++--
 pkg/server/mrt.go     |  2 +-
 pkg/server/server.go  | 24 +++++++++++++++++++-----
 pkg/server/zclient.go |  2 +-
 4 files changed, 23 insertions(+), 9 deletions(-)

diff --git a/pkg/server/bmp.go b/pkg/server/bmp.go
index 22be03aef..839dd7862 100644
--- a/pkg/server/bmp.go
+++ b/pkg/server/bmp.go
@@ -122,10 +122,10 @@ func (b *bmpClient) loop() {
 				).Warn("both option for route-monitoring-policy is obsoleted")
 			}
 			if b.c.RouteMonitoringPolicy == config.BMP_ROUTE_MONITORING_POLICY_TYPE_PRE_POLICY || b.c.RouteMonitoringPolicy == config.BMP_ROUTE_MONITORING_POLICY_TYPE_ALL {
-				ops = append(ops, watchUpdate(true))
+				ops = append(ops, watchUpdate(true, ""))
 			}
 			if b.c.RouteMonitoringPolicy == config.BMP_ROUTE_MONITORING_POLICY_TYPE_POST_POLICY || b.c.RouteMonitoringPolicy == config.BMP_ROUTE_MONITORING_POLICY_TYPE_ALL {
-				ops = append(ops, watchPostUpdate(true))
+				ops = append(ops, watchPostUpdate(true, ""))
 			}
 			if b.c.RouteMonitoringPolicy == config.BMP_ROUTE_MONITORING_POLICY_TYPE_LOCAL_RIB || b.c.RouteMonitoringPolicy == config.BMP_ROUTE_MONITORING_POLICY_TYPE_ALL {
 				ops = append(ops, watchBestPath(true))
diff --git a/pkg/server/mrt.go b/pkg/server/mrt.go
index 773f0ea5c..6ce64c14f 100644
--- a/pkg/server/mrt.go
+++ b/pkg/server/mrt.go
@@ -51,7 +51,7 @@ func (m *mrtWriter) loop() error {
 	ops := []watchOption{}
 	switch m.c.DumpType {
 	case config.MRT_TYPE_UPDATES:
-		ops = append(ops, watchUpdate(false))
+		ops = append(ops, watchUpdate(false, ""))
 	case config.MRT_TYPE_TABLE:
 		if len(m.c.TableName) > 0 {
 			ops = append(ops, watchTableName(m.c.TableName))
diff --git a/pkg/server/server.go b/pkg/server/server.go
index 274f7177f..ee58585fd 100644
--- a/pkg/server/server.go
+++ b/pkg/server/server.go
@@ -3689,9 +3689,9 @@ func (s *BgpServer) MonitorTable(ctx context.Context, r *api.MonitorTableRequest
 			return s.watch(watchBestPath(r.Current)), nil
 		case api.TableType_ADJ_IN:
 			if r.PostPolicy {
-				return s.watch(watchPostUpdate(r.Current)), nil
+				return s.watch(watchPostUpdate(r.Current, r.Name)), nil
 			}
-			return s.watch(watchUpdate(r.Current)), nil
+			return s.watch(watchUpdate(r.Current, r.Name)), nil
 		default:
 			return nil, fmt.Errorf("unsupported resource type: %v", r.TableType)
 		}
@@ -3731,6 +3731,9 @@ func (s *BgpServer) MonitorTable(ctx context.Context, r *api.MonitorTableRequest
 					if path == nil || (r.Family != nil && family != path.GetRouteFamily()) {
 						continue
 					}
+					if len(r.Name) > 0 && r.Name != path.GetSource().Address.String() {
+						continue
+					}
 					select {
 					case <-ctx.Done():
 						return
@@ -3879,6 +3882,7 @@ type watchOptions struct {
 	initPeerState  bool
 	tableName      string
 	recvMessage    bool
+	peerAddress    string
 }
 
 type watchOption func(*watchOptions)
@@ -3892,21 +3896,23 @@ func watchBestPath(current bool) watchOption {
 	}
 }
 
-func watchUpdate(current bool) watchOption {
+func watchUpdate(current bool, peerAddress string) watchOption {
 	return func(o *watchOptions) {
 		o.preUpdate = true
 		if current {
 			o.initUpdate = true
 		}
+		o.peerAddress = peerAddress
 	}
 }
 
-func watchPostUpdate(current bool) watchOption {
+func watchPostUpdate(current bool, peerAddress string) watchOption {
 	return func(o *watchOptions) {
 		o.postUpdate = true
 		if current {
 			o.initPostUpdate = true
 		}
+		o.peerAddress = peerAddress
 	}
 }
 
@@ -4088,10 +4094,14 @@ func (s *BgpServer) watch(opts ...watchOption) (w *watcher) {
 			for _, peer := range s.neighborMap {
 				peer.fsm.lock.RLock()
 				notEstablished := peer.fsm.state != bgp.BGP_FSM_ESTABLISHED
+				peerAddress := peer.fsm.peerInfo.Address.String()
 				peer.fsm.lock.RUnlock()
 				if notEstablished {
 					continue
 				}
+				if len(w.opts.peerAddress) > 0 && w.opts.peerAddress != peerAddress {
+					continue
+				}
 				configNeighbor := w.s.toConfig(peer, false)
 				for _, rf := range peer.configuredRFlist() {
 					peer.fsm.lock.RLock()
@@ -4146,9 +4156,13 @@ func (s *BgpServer) watch(opts ...watchOption) (w *watcher) {
 				for peerInfo, paths := range pathsByPeer {
 					// create copy which can be access to without mutex
 					var configNeighbor *config.Neighbor
-					if peer, ok := s.neighborMap[peerInfo.Address.String()]; ok {
+					peerAddress := peerInfo.Address.String()
+					if peer, ok := s.neighborMap[peerAddress]; ok {
 						configNeighbor = w.s.toConfig(peer, false)
 					}
+					if w.opts.peerAddress != "" && w.opts.peerAddress != peerAddress {
+						continue
+					}
 
 					w.notify(&watchEventUpdate{
 						PeerAS:      peerInfo.AS,
diff --git a/pkg/server/zclient.go b/pkg/server/zclient.go
index 5dac2a30a..645cb0f4c 100644
--- a/pkg/server/zclient.go
+++ b/pkg/server/zclient.go
@@ -364,7 +364,7 @@ func (z *zebraClient) updatePathByNexthopCache(paths []*table.Path) {
 func (z *zebraClient) loop() {
 	w := z.server.watch([]watchOption{
 		watchBestPath(true),
-		watchPostUpdate(true),
+		watchPostUpdate(true, ""),
 	}...)
 	defer w.Stop()