Merge pull request #10897 from opensourcerouting/safi-nht

zebra,staticd,*: SAFI_MULTICAST NHT groundwork

14 files changed, 161 insertions, 127 deletions

diff --git a/bgpd/bgp_nht.c b/bgpd/bgp_nht.c
index d768733545..d3ebc0e6a2 100644
--- a/bgpd/bgp_nht.c
+++ b/bgpd/bgp_nht.c
@@ -889,8 +889,9 @@ static void sendmsg_zebra_rnh(struct bgp_nexthop_cache *bnc, int command)
 			   zserv_command_string(command), &bnc->prefix,
 			   bnc->bgp->name_pretty);
 
-	ret = zclient_send_rnh(zclient, command, &bnc->prefix, exact_match,
-			       resolve_via_default, bnc->bgp->vrf_id);
+	ret = zclient_send_rnh(zclient, command, &bnc->prefix, SAFI_UNICAST,
+			       exact_match, resolve_via_default,
+			       bnc->bgp->vrf_id);
 	if (ret == ZCLIENT_SEND_FAILURE) {
 		flog_warn(EC_BGP_ZEBRA_SEND,
 			  "sendmsg_nexthop: zclient_send_message() failed");
diff --git a/lib/zclient.c b/lib/zclient.c
index f6c5a8af08..0c34214151 100644
--- a/lib/zclient.c
+++ b/lib/zclient.c
@@ -768,9 +768,9 @@ static void zclient_connect(struct thread *t)
 }
 
 enum zclient_send_status zclient_send_rnh(struct zclient *zclient, int command,
-					  const struct prefix *p,
-					  bool connected,
-					  bool resolve_via_def, vrf_id_t vrf_id)
+					  const struct prefix *p, safi_t safi,
+					  bool connected, bool resolve_via_def,
+					  vrf_id_t vrf_id)
 {
 	struct stream *s;
 
@@ -779,7 +779,7 @@ enum zclient_send_status zclient_send_rnh(struct zclient *zclient, int command,
 	zclient_create_header(s, command, vrf_id);
 	stream_putc(s, (connected) ? 1 : 0);
 	stream_putc(s, (resolve_via_def) ? 1 : 0);
-	stream_putw(s, SAFI_UNICAST);
+	stream_putw(s, safi);
 	stream_putw(s, PREFIX_FAMILY(p));
 	stream_putc(s, p->prefixlen);
 	switch (PREFIX_FAMILY(p)) {
diff --git a/lib/zclient.h b/lib/zclient.h
index 092754f602..7e1283d830 100644
--- a/lib/zclient.h
+++ b/lib/zclient.h
@@ -1071,7 +1071,8 @@ extern enum zclient_send_status zclient_route_send(uint8_t, struct zclient *,
 						   struct zapi_route *);
 extern enum zclient_send_status
 zclient_send_rnh(struct zclient *zclient, int command, const struct prefix *p,
-		 bool connected, bool resolve_via_default, vrf_id_t vrf_id);
+		 safi_t safi, bool connected, bool resolve_via_default,
+		 vrf_id_t vrf_id);
 int zapi_nexthop_encode(struct stream *s, const struct zapi_nexthop *api_nh,
 			uint32_t api_flags, uint32_t api_message);
 extern int zapi_route_encode(uint8_t, struct stream *, struct zapi_route *);
diff --git a/ospf6d/ospf6_zebra.c b/ospf6d/ospf6_zebra.c
index 4b37c5bc20..18b1959b9b 100644
--- a/ospf6d/ospf6_zebra.c
+++ b/ospf6d/ospf6_zebra.c
@@ -155,8 +155,8 @@ void ospf6_zebra_import_default_route(struct ospf6 *ospf6, bool unreg)
 			   zserv_command_string(command), &prefix,
 			   ospf6->vrf_id);
 
-	if (zclient_send_rnh(zclient, command, &prefix, false, true,
-			     ospf6->vrf_id)
+	if (zclient_send_rnh(zclient, command, &prefix, SAFI_UNICAST, false,
+			     true, ospf6->vrf_id)
 	    == ZCLIENT_SEND_FAILURE)
 		flog_err(EC_LIB_ZAPI_SOCKET, "%s: zclient_send_rnh() failed",
 			 __func__);
diff --git a/pbrd/pbr_zebra.c b/pbrd/pbr_zebra.c
index f992588729..da4e3e1bc0 100644
--- a/pbrd/pbr_zebra.c
+++ b/pbrd/pbr_zebra.c
@@ -485,7 +485,8 @@ void pbr_send_rnh(struct nexthop *nhop, bool reg)
 		break;
 	}
 
-	if (zclient_send_rnh(zclient, command, &p, false, false, nhop->vrf_id)
+	if (zclient_send_rnh(zclient, command, &p, SAFI_UNICAST, false, false,
+			     nhop->vrf_id)
 	    == ZCLIENT_SEND_FAILURE) {
 		zlog_warn("%s: Failure to send nexthop to zebra", __func__);
 	}
diff --git a/pimd/pim_nht.c b/pimd/pim_nht.c
index 753857005a..94dcfb8265 100644
--- a/pimd/pim_nht.c
+++ b/pimd/pim_nht.c
@@ -55,7 +55,7 @@ void pim_sendmsg_zebra_rnh(struct pim_instance *pim, struct zclient *zclient,
 	int ret;
 
 	p = &(pnc->rpf.rpf_addr);
-	ret = zclient_send_rnh(zclient, command, p, false, false,
+	ret = zclient_send_rnh(zclient, command, p, SAFI_UNICAST, false, false,
 			       pim->vrf->vrf_id);
 	if (ret == ZCLIENT_SEND_FAILURE)
 		zlog_warn("sendmsg_nexthop: zclient_send_message() failed");
diff --git a/sharpd/sharp_zebra.c b/sharpd/sharp_zebra.c
index 52364bff41..b40c2c6228 100644
--- a/sharpd/sharp_zebra.c
+++ b/sharpd/sharp_zebra.c
@@ -631,7 +631,8 @@ void sharp_zebra_nexthop_watch(struct prefix *p, vrf_id_t vrf_id, bool import,
 	if (!watch)
 		command = ZEBRA_NEXTHOP_UNREGISTER;
 
-	if (zclient_send_rnh(zclient, command, p, connected, false, vrf_id)
+	if (zclient_send_rnh(zclient, command, p, SAFI_UNICAST, connected,
+			     false, vrf_id)
 	    == ZCLIENT_SEND_FAILURE)
 		zlog_warn("%s: Failure to send nexthop to zebra", __func__);
 }
diff --git a/staticd/static_nht.c b/staticd/static_nht.c
index e1d6ba15d1..1d87a83c91 100644
--- a/staticd/static_nht.c
+++ b/staticd/static_nht.c
@@ -101,18 +101,15 @@ static void static_nht_update_safi(struct prefix *sp, struct prefix *nhp,
 	}
 }
 
-void static_nht_update(struct prefix *sp, struct prefix *nhp,
-		       uint32_t nh_num, afi_t afi, vrf_id_t nh_vrf_id)
+void static_nht_update(struct prefix *sp, struct prefix *nhp, uint32_t nh_num,
+		       afi_t afi, safi_t safi, vrf_id_t nh_vrf_id)
 {
 
 	struct vrf *vrf;
 
-	RB_FOREACH (vrf, vrf_name_head, &vrfs_by_name) {
-		static_nht_update_safi(sp, nhp, nh_num, afi, SAFI_UNICAST,
-				       vrf, nh_vrf_id);
-		static_nht_update_safi(sp, nhp, nh_num, afi, SAFI_MULTICAST,
-				       vrf, nh_vrf_id);
-	}
+	RB_FOREACH (vrf, vrf_name_head, &vrfs_by_name)
+		static_nht_update_safi(sp, nhp, nh_num, afi, safi, vrf,
+				       nh_vrf_id);
 }
 
 static void static_nht_reset_start_safi(struct prefix *nhp, afi_t afi,
@@ -166,16 +163,13 @@ static void static_nht_reset_start_safi(struct prefix *nhp, afi_t afi,
 	}
 }
 
-void static_nht_reset_start(struct prefix *nhp, afi_t afi, vrf_id_t nh_vrf_id)
+void static_nht_reset_start(struct prefix *nhp, afi_t afi, safi_t safi,
+			    vrf_id_t nh_vrf_id)
 {
 	struct vrf *vrf;
 
-	RB_FOREACH (vrf, vrf_name_head, &vrfs_by_name) {
-		static_nht_reset_start_safi(nhp, afi, SAFI_UNICAST,
-					    vrf, nh_vrf_id);
-		static_nht_reset_start_safi(nhp, afi, SAFI_MULTICAST,
-					    vrf, nh_vrf_id);
-	}
+	RB_FOREACH (vrf, vrf_name_head, &vrfs_by_name)
+		static_nht_reset_start_safi(nhp, afi, safi, vrf, nh_vrf_id);
 }
 
 static void static_nht_mark_state_safi(struct prefix *sp, afi_t afi,
@@ -212,7 +206,7 @@ static void static_nht_mark_state_safi(struct prefix *sp, afi_t afi,
 	route_unlock_node(rn);
 }
 
-void static_nht_mark_state(struct prefix *sp, vrf_id_t vrf_id,
+void static_nht_mark_state(struct prefix *sp, safi_t safi, vrf_id_t vrf_id,
 			   enum static_install_states state)
 {
 	struct vrf *vrf;
@@ -226,6 +220,5 @@ void static_nht_mark_state(struct prefix *sp, vrf_id_t vrf_id,
 	if (!vrf || !vrf->info)
 		return;
 
-	static_nht_mark_state_safi(sp, afi, SAFI_UNICAST, vrf, state);
-	static_nht_mark_state_safi(sp, afi, SAFI_MULTICAST, vrf, state);
+	static_nht_mark_state_safi(sp, afi, safi, vrf, state);
 }
diff --git a/staticd/static_nht.h b/staticd/static_nht.h
index 08dba2ebb5..c29acc32ef 100644
--- a/staticd/static_nht.h
+++ b/staticd/static_nht.h
@@ -37,19 +37,21 @@ extern "C" {
  * vrf_id -> The vrf the nexthop is in.
  */
 extern void static_nht_update(struct prefix *sp, struct prefix *nhp,
-			      uint32_t nh_num, afi_t afi, vrf_id_t vrf_id);
+			      uint32_t nh_num, afi_t afi, safi_t safi,
+			      vrf_id_t vrf_id);
 
 /*
  * For the given tracked nexthop, nhp, mark all routes that use
  * this route as in starting state again.
  */
-extern void static_nht_reset_start(struct prefix *nhp, afi_t afi,
+extern void static_nht_reset_start(struct prefix *nhp, afi_t afi, safi_t safi,
 				   vrf_id_t nh_vrf_id);
 
 /*
  * For the given prefix, sp, mark it as in a particular state
  */
-extern void static_nht_mark_state(struct prefix *sp, vrf_id_t vrf_id,
+extern void static_nht_mark_state(struct prefix *sp, safi_t safi,
+				  vrf_id_t vrf_id,
 				  enum static_install_states state);
 
 /*
diff --git a/staticd/static_zebra.c b/staticd/static_zebra.c
index bd293edebc..af153b4bc3 100644
--- a/staticd/static_zebra.c
+++ b/staticd/static_zebra.c
@@ -44,9 +44,47 @@
 #include "static_vty.h"
 #include "static_debug.h"
 
+DEFINE_MTYPE_STATIC(STATIC, STATIC_NHT_DATA, "Static Nexthop tracking data");
+PREDECL_HASH(static_nht_hash);
+
+struct static_nht_data {
+	struct static_nht_hash_item itm;
+
+	struct prefix nh;
+	safi_t safi;
+
+	vrf_id_t nh_vrf_id;
+
+	uint32_t refcount;
+	uint8_t nh_num;
+};
+
+static int static_nht_data_cmp(const struct static_nht_data *nhtd1,
+			       const struct static_nht_data *nhtd2)
+{
+	if (nhtd1->nh_vrf_id != nhtd2->nh_vrf_id)
+		return numcmp(nhtd1->nh_vrf_id, nhtd2->nh_vrf_id);
+	if (nhtd1->safi != nhtd2->safi)
+		return numcmp(nhtd1->safi, nhtd2->safi);
+
+	return prefix_cmp(&nhtd1->nh, &nhtd2->nh);
+}
+
+static unsigned int static_nht_data_hash(const struct static_nht_data *nhtd)
+{
+	unsigned int key = 0;
+
+	key = prefix_hash_key(&nhtd->nh);
+	return jhash_2words(nhtd->nh_vrf_id, nhtd->safi, key);
+}
+
+DECLARE_HASH(static_nht_hash, struct static_nht_data, itm, static_nht_data_cmp,
+	     static_nht_data_hash);
+
+static struct static_nht_hash_head static_nht_hash[1];
+
 /* Zebra structure to hold current status. */
 struct zclient *zclient;
-static struct hash *static_nht_hash;
 uint32_t zebra_ecmp_count = MULTIPATH_NUM;
 
 /* Inteface addition message from zebra. */
@@ -104,31 +142,32 @@ static int route_notify_owner(ZAPI_CALLBACK_ARGS)
 	struct prefix p;
 	enum zapi_route_notify_owner note;
 	uint32_t table_id;
+	safi_t safi;
 
-	if (!zapi_route_notify_decode(zclient->ibuf, &p, &table_id, &note,
-				      NULL, NULL))
+	if (!zapi_route_notify_decode(zclient->ibuf, &p, &table_id, &note, NULL,
+				      &safi))
 		return -1;
 
 	switch (note) {
 	case ZAPI_ROUTE_FAIL_INSTALL:
-		static_nht_mark_state(&p, vrf_id, STATIC_NOT_INSTALLED);
+		static_nht_mark_state(&p, safi, vrf_id, STATIC_NOT_INSTALLED);
 		zlog_warn("%s: Route %pFX failed to install for table: %u",
 			  __func__, &p, table_id);
 		break;
 	case ZAPI_ROUTE_BETTER_ADMIN_WON:
-		static_nht_mark_state(&p, vrf_id, STATIC_NOT_INSTALLED);
+		static_nht_mark_state(&p, safi, vrf_id, STATIC_NOT_INSTALLED);
 		zlog_warn(
 			"%s: Route %pFX over-ridden by better route for table: %u",
 			__func__, &p, table_id);
 		break;
 	case ZAPI_ROUTE_INSTALLED:
-		static_nht_mark_state(&p, vrf_id, STATIC_INSTALLED);
+		static_nht_mark_state(&p, safi, vrf_id, STATIC_INSTALLED);
 		break;
 	case ZAPI_ROUTE_REMOVED:
-		static_nht_mark_state(&p, vrf_id, STATIC_NOT_INSTALLED);
+		static_nht_mark_state(&p, safi, vrf_id, STATIC_NOT_INSTALLED);
 		break;
 	case ZAPI_ROUTE_REMOVE_FAIL:
-		static_nht_mark_state(&p, vrf_id, STATIC_INSTALLED);
+		static_nht_mark_state(&p, safi, vrf_id, STATIC_INSTALLED);
 		zlog_warn("%s: Route %pFX failure to remove for table: %u",
 			  __func__, &p, table_id);
 		break;
@@ -141,15 +180,6 @@ static void zebra_connected(struct zclient *zclient)
 	zclient_send_reg_requests(zclient, VRF_DEFAULT);
 }
 
-struct static_nht_data {
-	struct prefix *nh;
-
-	vrf_id_t nh_vrf_id;
-
-	uint32_t refcount;
-	uint8_t nh_num;
-};
-
 /* API to check whether the configured nexthop address is
  * one of its local connected address or not.
  */
@@ -187,17 +217,19 @@ static int static_zebra_nexthop_update(ZAPI_CALLBACK_ARGS)
 	}
 
 	memset(&lookup, 0, sizeof(lookup));
-	lookup.nh = &matched;
+	lookup.nh = matched;
 	lookup.nh_vrf_id = vrf_id;
+	lookup.safi = nhr.safi;
 
-	nhtd = hash_lookup(static_nht_hash, &lookup);
+	nhtd = static_nht_hash_find(static_nht_hash, &lookup);
 
 	if (nhtd) {
 		nhtd->nh_num = nhr.nexthop_num;
 
-		static_nht_reset_start(&matched, afi, nhtd->nh_vrf_id);
+		static_nht_reset_start(&matched, afi, nhr.safi,
+				       nhtd->nh_vrf_id);
 		static_nht_update(NULL, &matched, nhr.nexthop_num, afi,
-				  nhtd->nh_vrf_id);
+				  nhr.safi, nhtd->nh_vrf_id);
 	} else
 		zlog_err("No nhtd?");
 
@@ -210,55 +242,50 @@ static void static_zebra_capabilities(struct zclient_capabilities *cap)
 	zebra_ecmp_count = cap->ecmp;
 }
 
-static unsigned int static_nht_hash_key(const void *data)
+static struct static_nht_data *
+static_nht_hash_getref(const struct static_nht_data *ref)
 {
-	const struct static_nht_data *nhtd = data;
-	unsigned int key = 0;
+	struct static_nht_data *nhtd;
 
-	key = prefix_hash_key(nhtd->nh);
-	return jhash_1word(nhtd->nh_vrf_id, key);
-}
+	nhtd = static_nht_hash_find(static_nht_hash, ref);
+	if (!nhtd) {
+		nhtd = XCALLOC(MTYPE_STATIC_NHT_DATA, sizeof(*nhtd));
 
-static bool static_nht_hash_cmp(const void *d1, const void *d2)
-{
-	const struct static_nht_data *nhtd1 = d1;
-	const struct static_nht_data *nhtd2 = d2;
+		prefix_copy(&nhtd->nh, &ref->nh);
+		nhtd->nh_vrf_id = ref->nh_vrf_id;
+		nhtd->safi = ref->safi;
 
-	if (nhtd1->nh_vrf_id != nhtd2->nh_vrf_id)
-		return false;
+		static_nht_hash_add(static_nht_hash, nhtd);
+	}
 
-	return prefix_same(nhtd1->nh, nhtd2->nh);
+	nhtd->refcount++;
+	return nhtd;
 }
 
-static void *static_nht_hash_alloc(void *data)
+static bool static_nht_hash_decref(struct static_nht_data *nhtd)
 {
-	struct static_nht_data *copy = data;
-	struct static_nht_data *new;
-
-	new = XMALLOC(MTYPE_TMP, sizeof(*new));
+	if (--nhtd->refcount > 0)
+		return true;
 
-	new->nh = prefix_new();
-	prefix_copy(new->nh, copy->nh);
-	new->refcount = 0;
-	new->nh_num = 0;
-	new->nh_vrf_id = copy->nh_vrf_id;
-
-	return new;
+	static_nht_hash_del(static_nht_hash, nhtd);
+	XFREE(MTYPE_STATIC_NHT_DATA, nhtd);
+	return false;
 }
 
-static void static_nht_hash_free(void *data)
+static void static_nht_hash_clear(void)
 {
-	struct static_nht_data *nhtd = data;
+	struct static_nht_data *nhtd;
 
-	prefix_free(&nhtd->nh);
-	XFREE(MTYPE_TMP, nhtd);
+	while ((nhtd = static_nht_hash_pop(static_nht_hash)))
+		XFREE(MTYPE_STATIC_NHT_DATA, nhtd);
 }
 
 void static_zebra_nht_register(struct static_nexthop *nh, bool reg)
 {
 	struct static_path *pn = nh->pn;
 	struct route_node *rn = pn->rn;
-	struct static_nht_data *nhtd, lookup;
+	struct static_route_info *si = static_route_info_from_rnode(rn);
+	struct static_nht_data lookup;
 	uint32_t cmd;
 	struct prefix p;
 	afi_t afi = AFI_IP;
@@ -294,44 +321,42 @@ void static_zebra_nht_register(struct static_nexthop *nh, bool reg)
 	}
 
 	memset(&lookup, 0, sizeof(lookup));
-	lookup.nh = &p;
+	lookup.nh = p;
 	lookup.nh_vrf_id = nh->nh_vrf_id;
+	lookup.safi = si->safi;
 
 	nh->nh_registered = reg;
 
 	if (reg) {
-		nhtd = hash_get(static_nht_hash, &lookup,
-				static_nht_hash_alloc);
-		nhtd->refcount++;
+		struct static_nht_data *nhtd;
+
+		nhtd = static_nht_hash_getref(&lookup);
 
 		if (nhtd->refcount > 1) {
 			DEBUGD(&static_dbg_route,
 			       "Already registered nexthop(%pFX) for %pRN %d",
 			       &p, rn, nhtd->nh_num);
 			if (nhtd->nh_num)
-				static_nht_update(&rn->p, nhtd->nh,
-						  nhtd->nh_num, afi,
+				static_nht_update(&rn->p, &nhtd->nh,
+						  nhtd->nh_num, afi, si->safi,
 						  nh->nh_vrf_id);
 			return;
 		}
 	} else {
-		nhtd = hash_lookup(static_nht_hash, &lookup);
+		struct static_nht_data *nhtd;
+
+		nhtd = static_nht_hash_find(static_nht_hash, &lookup);
 		if (!nhtd)
 			return;
-
-		nhtd->refcount--;
-		if (nhtd->refcount >= 1)
+		if (static_nht_hash_decref(nhtd))
 			return;
-
-		hash_release(static_nht_hash, nhtd);
-		static_nht_hash_free(nhtd);
 	}
 
 	DEBUGD(&static_dbg_route, "%s nexthop(%pFX) for %pRN",
 	       reg ? "Registering" : "Unregistering", &p, rn);
 
-	if (zclient_send_rnh(zclient, cmd, &p, false, false, nh->nh_vrf_id)
-	    == ZCLIENT_SEND_FAILURE)
+	if (zclient_send_rnh(zclient, cmd, &p, si->safi, false, false,
+			     nh->nh_vrf_id) == ZCLIENT_SEND_FAILURE)
 		zlog_warn("%s: Failure to send nexthop to zebra", __func__);
 }
 /*
@@ -342,6 +367,7 @@ int static_zebra_nh_update(struct static_nexthop *nh)
 {
 	struct static_path *pn = nh->pn;
 	struct route_node *rn = pn->rn;
+	struct static_route_info *si = static_route_info_from_rnode(rn);
 	struct static_nht_data *nhtd, lookup = {};
 	struct prefix p = {};
 	afi_t afi = AFI_IP;
@@ -369,14 +395,15 @@ int static_zebra_nh_update(struct static_nexthop *nh)
 		break;
 	}
 
-	lookup.nh = &p;
+	lookup.nh = p;
 	lookup.nh_vrf_id = nh->nh_vrf_id;
+	lookup.safi = si->safi;
 
-	nhtd = hash_lookup(static_nht_hash, &lookup);
+	nhtd = static_nht_hash_find(static_nht_hash, &lookup);
 	if (nhtd && nhtd->nh_num) {
 		nh->state = STATIC_START;
-		static_nht_update(&rn->p, nhtd->nh, nhtd->nh_num, afi,
-				  nh->nh_vrf_id);
+		static_nht_update(&rn->p, &nhtd->nh, nhtd->nh_num, afi,
+				  si->safi, nh->nh_vrf_id);
 		return 1;
 	}
 	return 0;
@@ -531,14 +558,15 @@ void static_zebra_init(void)
 	zclient->zebra_capabilities = static_zebra_capabilities;
 	zclient->zebra_connected = zebra_connected;
 
-	static_nht_hash = hash_create(static_nht_hash_key,
-				      static_nht_hash_cmp,
-				      "Static Nexthop Tracking hash");
+	static_nht_hash_init(static_nht_hash);
 }
 
 /* static_zebra_stop used by tests/lib/test_grpc.cpp */
 void static_zebra_stop(void)
 {
+	static_nht_hash_clear();
+	static_nht_hash_fini(static_nht_hash);
+
 	if (!zclient)
 		return;
 	zclient_stop(zclient);
diff --git a/zebra/zapi_msg.c b/zebra/zapi_msg.c
index e94aee5c1a..fd475e4cee 100644
--- a/zebra/zapi_msg.c
+++ b/zebra/zapi_msg.c
@@ -1215,7 +1215,7 @@ static void zread_rnh_register(ZAPI_HANDLER_ARGS)
 				p.family);
 			return;
 		}
-		rnh = zebra_add_rnh(&p, zvrf_id(zvrf), &exist);
+		rnh = zebra_add_rnh(&p, zvrf_id(zvrf), safi, &exist);
 		if (!rnh)
 			return;
 
diff --git a/zebra/zebra_rnh.c b/zebra/zebra_rnh.c
index 4d5336120d..7934a9d206 100644
--- a/zebra/zebra_rnh.c
+++ b/zebra/zebra_rnh.c
@@ -132,13 +132,13 @@ static void zebra_rnh_store_in_routing_table(struct rnh *rnh)
 	route_unlock_node(rn);
 }
 
-struct rnh *zebra_add_rnh(struct prefix *p, vrf_id_t vrfid, bool *exists)
+struct rnh *zebra_add_rnh(struct prefix *p, vrf_id_t vrfid, safi_t safi,
+			  bool *exists)
 {
 	struct route_table *table;
 	struct route_node *rn;
 	struct rnh *rnh = NULL;
 	afi_t afi = family2afi(p->family);
-	safi_t safi = SAFI_UNICAST;
 
 	if (IS_ZEBRA_DEBUG_NHT) {
 		struct vrf *vrf = vrf_lookup_by_id(vrfid);
@@ -345,7 +345,7 @@ void zebra_register_rnh_pseudowire(vrf_id_t vrf_id, struct zebra_pw *pw,
 		return;
 
 	addr2hostprefix(pw->af, &pw->nexthop, &nh);
-	rnh = zebra_add_rnh(&nh, vrf_id, &exists);
+	rnh = zebra_add_rnh(&nh, vrf_id, SAFI_UNICAST, &exists);
 	if (!rnh)
 		return;
 
@@ -768,7 +768,7 @@ static void zebra_rnh_clear_nhc_flag(struct zebra_vrf *zvrf, afi_t afi,
  * of a particular VRF and address-family or a specific prefix.
  */
 void zebra_evaluate_rnh(struct zebra_vrf *zvrf, afi_t afi, int force,
-			struct prefix *p, safi_t safi)
+			const struct prefix *p, safi_t safi)
 {
 	struct route_table *rnh_table;
 	struct route_node *nrn;
@@ -802,13 +802,13 @@ void zebra_evaluate_rnh(struct zebra_vrf *zvrf, afi_t afi, int force,
 	}
 }
 
-void zebra_print_rnh_table(vrf_id_t vrfid, afi_t afi, struct vty *vty,
-			   struct prefix *p)
+void zebra_print_rnh_table(vrf_id_t vrfid, afi_t afi, safi_t safi,
+			   struct vty *vty, const struct prefix *p)
 {
 	struct route_table *table;
 	struct route_node *rn;
 
-	table = get_rnh_table(vrfid, afi, SAFI_UNICAST);
+	table = get_rnh_table(vrfid, afi, safi);
 	if (!table) {
 		if (IS_ZEBRA_DEBUG_NHT)
 			zlog_debug("print_rnhs: rnh table not found");
@@ -1337,7 +1337,7 @@ static void print_rnh(struct route_node *rn, struct vty *vty)
 	vty_out(vty, "\n");
 }
 
-static int zebra_cleanup_rnh_client(vrf_id_t vrf_id, afi_t afi,
+static int zebra_cleanup_rnh_client(vrf_id_t vrf_id, afi_t afi, safi_t safi,
 				    struct zserv *client)
 {
 	struct route_table *ntable;
@@ -1352,7 +1352,7 @@ static int zebra_cleanup_rnh_client(vrf_id_t vrf_id, afi_t afi,
 			   zebra_route_string(client->proto), afi2str(afi));
 	}
 
-	ntable = get_rnh_table(vrf_id, afi, SAFI_UNICAST);
+	ntable = get_rnh_table(vrf_id, afi, safi);
 	if (!ntable) {
 		zlog_debug("cleanup_rnh_client: rnh table not found");
 		return -1;
@@ -1377,9 +1377,14 @@ static int zebra_client_cleanup_rnh(struct zserv *client)
 	RB_FOREACH (vrf, vrf_id_head, &vrfs_by_id) {
 		zvrf = vrf->info;
 		if (zvrf) {
-			zebra_cleanup_rnh_client(zvrf_id(zvrf), AFI_IP, client);
+			zebra_cleanup_rnh_client(zvrf_id(zvrf), AFI_IP,
+						 SAFI_UNICAST, client);
+			zebra_cleanup_rnh_client(zvrf_id(zvrf), AFI_IP,
+						 SAFI_MULTICAST, client);
+			zebra_cleanup_rnh_client(zvrf_id(zvrf), AFI_IP6,
+						 SAFI_UNICAST, client);
 			zebra_cleanup_rnh_client(zvrf_id(zvrf), AFI_IP6,
-						 client);
+						 SAFI_MULTICAST, client);
 		}
 	}
 
diff --git a/zebra/zebra_rnh.h b/zebra/zebra_rnh.h
index 27c016ebe6..70eda725c4 100644
--- a/zebra/zebra_rnh.h
+++ b/zebra/zebra_rnh.h
@@ -31,7 +31,7 @@ extern "C" {
 
 extern void zebra_rnh_init(void);
 
-extern struct rnh *zebra_add_rnh(struct prefix *p, vrf_id_t vrfid,
+extern struct rnh *zebra_add_rnh(struct prefix *p, vrf_id_t vrfid, safi_t safi,
 				 bool *exists);
 extern struct rnh *zebra_lookup_rnh(struct prefix *p, vrf_id_t vrfid,
 				    safi_t safi);
@@ -44,9 +44,9 @@ extern void zebra_register_rnh_pseudowire(vrf_id_t, struct zebra_pw *, bool *);
 extern void zebra_deregister_rnh_pseudowire(vrf_id_t, struct zebra_pw *);
 extern void zebra_remove_rnh_client(struct rnh *rnh, struct zserv *client);
 extern void zebra_evaluate_rnh(struct zebra_vrf *zvrf, afi_t afi, int force,
-			       struct prefix *p, safi_t safi);
-extern void zebra_print_rnh_table(vrf_id_t vrfid, afi_t afi, struct vty *vty,
-				  struct prefix *p);
+			       const struct prefix *p, safi_t safi);
+extern void zebra_print_rnh_table(vrf_id_t vrfid, afi_t afi, safi_t safi,
+				  struct vty *vty, const struct prefix *p);
 
 extern int rnh_resolve_via_default(struct zebra_vrf *zvrf, int family);
 
diff --git a/zebra/zebra_vty.c b/zebra/zebra_vty.c
index 32bbfd6654..22c65e3c0c 100644
--- a/zebra/zebra_vty.c
+++ b/zebra/zebra_vty.c
@@ -1361,7 +1361,7 @@ static int do_show_ip_route(struct vty *vty, const char *vrf_name, afi_t afi,
 
 DEFPY (show_ip_nht,
        show_ip_nht_cmd,
-       "show <ip$ipv4|ipv6$ipv6> <nht|import-check>$type [<A.B.C.D|X:X::X:X>$addr|vrf NAME$vrf_name [<A.B.C.D|X:X::X:X>$addr]|vrf all$vrf_all]",
+       "show <ip$ipv4|ipv6$ipv6> <nht|import-check>$type [<A.B.C.D|X:X::X:X>$addr|vrf NAME$vrf_name [<A.B.C.D|X:X::X:X>$addr]|vrf all$vrf_all] [mrib$mrib]",
        SHOW_STR
        IP_STR
        IP6_STR
@@ -1372,11 +1372,13 @@ DEFPY (show_ip_nht,
        VRF_CMD_HELP_STR
        "IPv4 Address\n"
        "IPv6 Address\n"
-       VRF_ALL_CMD_HELP_STR)
+       VRF_ALL_CMD_HELP_STR
+       "Show Multicast (MRIB) NHT state\n")
 {
 	afi_t afi = ipv4 ? AFI_IP : AFI_IP6;
 	vrf_id_t vrf_id = VRF_DEFAULT;
 	struct prefix prefix, *p = NULL;
+	safi_t safi = mrib ? SAFI_MULTICAST : SAFI_UNICAST;
 
 	if (vrf_all) {
 		struct vrf *vrf;
@@ -1385,8 +1387,8 @@ DEFPY (show_ip_nht,
 		RB_FOREACH (vrf, vrf_name_head, &vrfs_by_name)
 			if ((zvrf = vrf->info) != NULL) {
 				vty_out(vty, "\nVRF %s:\n", zvrf_name(zvrf));
-				zebra_print_rnh_table(zvrf_id(zvrf), afi, vty,
-						      NULL);
+				zebra_print_rnh_table(zvrf_id(zvrf), afi, safi,
+						      vty, NULL);
 			}
 		return CMD_SUCCESS;
 	}
@@ -1400,7 +1402,7 @@ DEFPY (show_ip_nht,
 			return CMD_WARNING;
 	}
 
-	zebra_print_rnh_table(vrf_id, afi, vty, p);
+	zebra_print_rnh_table(vrf_id, afi, safi, vty, p);
 	return CMD_SUCCESS;
 }