3 files changed, 53 insertions, 0 deletions

diff --git a/zebra/zebra_nhg.c b/zebra/zebra_nhg.c
index d4d31d3ea2..407af2902d 100644
--- a/zebra/zebra_nhg.c
+++ b/zebra/zebra_nhg.c
@@ -2813,3 +2813,40 @@ struct nhg_hash_entry *zebra_nhg_proto_del(uint32_t id)
 	return nhe;
 }
 
+struct nhg_score_proto_iter {
+	int type;
+	unsigned long found;
+};
+
+static void zebra_nhg_score_proto_entry(struct hash_bucket *bucket, void *arg)
+{
+	struct nhg_hash_entry *nhe;
+	struct nhg_score_proto_iter *iter;
+
+	nhe = (struct nhg_hash_entry *)bucket->data;
+	iter = arg;
+
+	/* Needs to match type and outside zebra ID space */
+	if (nhe->type == iter->type && nhe->id >= ZEBRA_NHG_PROTO_LOWER) {
+		if (IS_ZEBRA_DEBUG_NHG_DETAIL)
+			zlog_debug(
+				"%s: found nhe %p (%u), vrf %d, type %s after client disconnect",
+				__func__, nhe, nhe->id, nhe->vrf_id,
+				zebra_route_string(nhe->type));
+
+		/* This should be the last ref if we remove client routes too */
+		zebra_nhg_decrement_ref(nhe);
+	}
+}
+
+/* Remove specific by proto NHGs */
+unsigned long zebra_nhg_score_proto(int type)
+{
+	struct nhg_score_proto_iter iter = {};
+
+	iter.type = type;
+
+	hash_iterate(zrouter.nhgs_id, zebra_nhg_score_proto_entry, &iter);
+
+	return iter.found;
+}
diff --git a/zebra/zebra_nhg.h b/zebra/zebra_nhg.h
index 16474aada4..8c33a91917 100644
--- a/zebra/zebra_nhg.h
+++ b/zebra/zebra_nhg.h
@@ -285,6 +285,14 @@ struct nhg_hash_entry *zebra_nhg_proto_add(uint32_t id, int type,
  */
 struct nhg_hash_entry *zebra_nhg_proto_del(uint32_t id);
 
+/*
+ * Remove specific by proto NHGs.
+ *
+ * Called after client disconnect.
+ *
+ */
+unsigned long zebra_nhg_score_proto(int type);
+
 /* Reference counter functions */
 extern void zebra_nhg_decrement_ref(struct nhg_hash_entry *nhe);
 extern void zebra_nhg_increment_ref(struct nhg_hash_entry *nhe);
diff --git a/zebra/zserv.c b/zebra/zserv.c
index 4c8656af0d..44f4641fcf 100644
--- a/zebra/zserv.c
+++ b/zebra/zserv.c
@@ -590,6 +590,7 @@ static void zserv_client_free(struct zserv *client)
 	/* Close file descriptor. */
 	if (client->sock) {
 		unsigned long nroutes;
+		unsigned long nnhgs;
 
 		close(client->sock);
 
@@ -600,6 +601,13 @@ static void zserv_client_free(struct zserv *client)
 				"client %d disconnected %lu %s routes removed from the rib",
 				client->sock, nroutes,
 				zebra_route_string(client->proto));
+
+			/* Not worrying about instance for now */
+			nnhgs = zebra_nhg_score_proto(client->proto);
+			zlog_notice(
+				"client %d disconnected %lu %s nhgs removed from the rib",
+				client->sock, nnhgs,
+				zebra_route_string(client->proto));
 		}
 		client->sock = -1;
 	}