6 files changed, 232 insertions, 182 deletions

diff --git a/zebra/interface.c b/zebra/interface.c
index 3eb0e6853..8fe7af3f2 100644
--- a/zebra/interface.c
+++ b/zebra/interface.c
@@ -53,7 +53,6 @@
 #include "zebra/zebra_errors.h"
 
 DEFINE_MTYPE_STATIC(ZEBRA, ZINFO, "Zebra Interface Information")
-DEFINE_MTYPE_STATIC(ZEBRA, NHE_CONNECTED, "Nexthops Connected")
 
 #define ZEBRA_PTM_SUPPORT
 
@@ -108,6 +107,17 @@ static void zebra_if_node_destroy(route_table_delegate_t *delegate,
 	route_node_destroy(delegate, table, node);
 }
 
+static void zebra_if_nhg_dependents_free(struct zebra_if *zebra_if)
+{
+	nhg_connected_head_free(&zebra_if->nhg_dependents);
+}
+
+static void zebra_if_nhg_dependents_init(struct zebra_if *zebra_if)
+{
+	nhg_connected_head_init(&zebra_if->nhg_dependents);
+}
+
+
 route_table_delegate_t zebra_if_table_delegate = {
 	.create_node = route_node_create,
 	.destroy_node = zebra_if_node_destroy};
@@ -122,8 +132,7 @@ static int if_zebra_new_hook(struct interface *ifp)
 	zebra_if->multicast = IF_ZEBRA_MULTICAST_UNSPEC;
 	zebra_if->shutdown = IF_ZEBRA_SHUTDOWN_OFF;
 
-	zebra_if->nhe_connected = list_new();
-	zebra_if->nhe_connected->del = (void (*)(void *))nhe_connected_free;
+	zebra_if_nhg_dependents_init(zebra_if);
 
 	zebra_ptm_if_init(zebra_if);
 
@@ -201,9 +210,10 @@ static int if_zebra_delete_hook(struct interface *ifp)
 		list_delete(&rtadv->AdvDNSSLList);
 #endif /* HAVE_RTADV */
 
-		list_delete(&zebra_if->nhe_connected);
+		zebra_if_nhg_dependents_free(zebra_if);
 
 		XFREE(MTYPE_TMP, zebra_if->desc);
+
 		THREAD_OFF(zebra_if->speed_update);
 
 		XFREE(MTYPE_ZINFO, zebra_if);
@@ -932,48 +942,48 @@ static void if_down_del_nbr_connected(struct interface *ifp)
 	}
 }
 
-/**
- * nhe_connected_add() - Add the nexthop entry to the interfaces connected list
- *
- * @ifp:	Interface to add to
- * @nhe:	Nexthop hash entry to add
- *
- * Return:	nhe_connected struct created and added
- */
-struct nhe_connected *nhe_connected_add(struct interface *ifp,
-					struct nhg_hash_entry *nhe)
+void if_nhg_dependents_add(struct interface *ifp, struct nhg_hash_entry *nhe)
 {
-	struct nhe_connected *if_nhec = NULL;
-	struct zebra_if *zif = (struct zebra_if *)ifp->info;
+	if (ifp->info) {
+		struct zebra_if *zif = (struct zebra_if *)ifp->info;
 
-	if_nhec = nhe_connected_new();
+		nhg_connected_head_add(&zif->nhg_dependents, nhe);
+	}
+}
 
-	/* Attach the nhe */
-	if_nhec->nhe = nhe;
+void if_nhg_dependents_del(struct interface *ifp, struct nhg_hash_entry *nhe)
+{
+	if (ifp->info) {
+		struct zebra_if *zif = (struct zebra_if *)ifp->info;
 
-	/* Add connected nexthop to the interface */
-	listnode_add(zif->nhe_connected, if_nhec);
-	return if_nhec;
+		nhg_connected_head_del(&zif->nhg_dependents, nhe);
+	}
 }
 
-/**
- * nhe_connected() - Allocate nhe connected structure
- *
- * Return:	Allocated nhe_connected structure
- */
-struct nhe_connected *nhe_connected_new(void)
+unsigned int if_nhg_dependents_count(const struct interface *ifp)
 {
-	return XCALLOC(MTYPE_NHE_CONNECTED, sizeof(struct nhe_connected));
+	if (ifp->info) {
+		struct zebra_if *zif = (struct zebra_if *)ifp->info;
+
+		return nhg_connected_head_count(&zif->nhg_dependents);
+	}
+
+	return 0;
+}
+
+
+bool if_nhg_dependents_is_empty(const struct interface *ifp)
+{
+	if (ifp->info) {
+		struct zebra_if *zif = (struct zebra_if *)ifp->info;
+
+		return nhg_connected_head_is_empty(&zif->nhg_dependents);
+	}
+
+	return false;
 }
 
-/**
- * nhe_connected_free() - Free nhe_connected structure
- *
- * @nhe_connected:	nhe_connected structure to free
- */
-void nhe_connected_free(struct nhe_connected *connected)
 {
-	XFREE(MTYPE_NHE_CONNECTED, connected);
 }
 
 /* Interface is up. */
@@ -1183,16 +1193,9 @@ static void nbr_connected_dump_vty(struct vty *vty,
 	vty_out(vty, "\n");
 }
 
-/**
- * nhe_connected_dump() - Dump nexthops connected to this interface to vty
- *
- * @vty:		Vty output
- * @nhe_connected:	List of connected nexthop hash entries
- */
-static void nhe_connected_dump_vty(struct vty *vty,
-				   struct nhe_connected *connected)
+static void nhg_dependent_dump_vty(struct vty *vty,
+				   struct nhg_connected *connected)
 {
-	/* Just outputing ID for now. */
 	vty_out(vty, "  (%u)", connected->nhe->id);
 }
 
@@ -1343,7 +1346,6 @@ static void if_dump_vty(struct vty *vty, struct interface *ifp)
 {
 	struct connected *connected;
 	struct nbr_connected *nbr_connected;
-	struct nhe_connected *nhe_connected;
 	struct listnode *node;
 	struct route_node *rn;
 	struct zebra_if *zebra_if;
@@ -1429,11 +1431,14 @@ static void if_dump_vty(struct vty *vty, struct interface *ifp)
 			connected_dump_vty(vty, connected);
 	}
 
-	if (listhead(zebra_if->nhe_connected)) {
+	if (!if_nhg_dependents_is_empty(ifp)) {
+		struct nhg_connected *rb_node_dep = NULL;
+
 		vty_out(vty, "  Nexthop IDs connected:");
-		for (ALL_LIST_ELEMENTS_RO(zebra_if->nhe_connected, node,
-					  nhe_connected))
-			nhe_connected_dump_vty(vty, nhe_connected);
+		RB_FOREACH (rb_node_dep, nhg_connected_head,
+			    &zebra_if->nhg_dependents) {
+			nhg_dependent_dump_vty(vty, rb_node_dep);
+		}
 		vty_out(vty, "\n");
 	}
 
diff --git a/zebra/interface.h b/zebra/interface.h
index 7dccaeacc..d5c1e1713 100644
--- a/zebra/interface.h
+++ b/zebra/interface.h
@@ -264,12 +264,6 @@ typedef enum {
 
 struct irdp_interface;
 
-/* Nexthop hash entry connected structure */
-struct nhe_connected {
-	/* Connected nexthop hash entry */
-	struct nhg_hash_entry *nhe;
-};
-
 /* `zebra' daemon local interface structure. */
 struct zebra_if {
 	/* Shutdown configuration. */
@@ -284,14 +278,14 @@ struct zebra_if {
 	/* Installed addresses chains tree. */
 	struct route_table *ipv4_subnets;
 
-	/* Nexthops pointed to it list */
+	/* Nexthops pointing to this interface */
 	/**
 	 * Any nexthop that we get should have an
 	 * interface. When an interface goes down,
 	 * we will use this list to update the nexthops
 	 * pointing to it with that info.
 	 */
-	struct list *nhe_connected;
+	struct nhg_connected_head nhg_dependents;
 
 	/* Information about up/down changes */
 	unsigned int up_count;
@@ -440,11 +434,13 @@ extern void zebra_if_update_link(struct interface *ifp, ifindex_t link_ifindex,
 extern void zebra_if_update_all_links(void);
 extern void zebra_if_set_protodown(struct interface *ifp, bool down);
 
-/* Nexthop connected list functions */
-struct nhe_connected *nhe_connected_add(struct interface *ifp,
-					struct nhg_hash_entry *nhe);
-struct nhe_connected *nhe_connected_new(void);
-void nhe_connected_free(struct nhe_connected *connected);
+/* Nexthop group connected functions */
+extern void if_nhg_dependents_add(struct interface *ifp,
+				  struct nhg_hash_entry *nhe);
+extern void if_nhg_dependents_del(struct interface *ifp,
+				  struct nhg_hash_entry *nhe);
+extern unsigned int if_nhg_dependents_count(const struct interface *ifp);
+extern bool if_nhg_dependents_is_empty(const struct interface *ifp);
 
 extern void vrf_add_update(struct vrf *vrfp);
 
diff --git a/zebra/rt_netlink.c b/zebra/rt_netlink.c
index bdf613cef..d2d9060b8 100644
--- a/zebra/rt_netlink.c
+++ b/zebra/rt_netlink.c
@@ -763,6 +763,7 @@ static int netlink_route_change_read_unicast(struct nlmsghdr *h, ns_id_t ns_id,
 				XFREE(MTYPE_RE, re);
 		}
 	} else {
+		// TODO: Use nhe_id here as well
 		if (!tb[RTA_MULTIPATH]) {
 			struct nexthop nh;
 			size_t sz = (afi == AFI_IP) ? 4 : 16;
@@ -2251,7 +2252,7 @@ static int netlink_nexthop_process_group(struct rtattr **tb,
 	zlog_debug("Nexthop group type: %d",
 		   *((uint16_t *)RTA_DATA(tb[NHA_GROUP_TYPE])));
 
-	zebra_nhg_connected_head_init(nhg_depends);
+	nhg_connected_head_init(nhg_depends);
 
 	for (int i = 0; i < count; i++) {
 		/* We do not care about nexthop_grp.weight at
@@ -2261,7 +2262,7 @@ static int netlink_nexthop_process_group(struct rtattr **tb,
 		 */
 		depend = zebra_nhg_lookup_id(n_grp[i].id);
 		if (depend) {
-			zebra_nhg_connected_head_add(nhg_depends, depend);
+			nhg_connected_head_add(nhg_depends, depend);
 			/*
 			 * If this is a nexthop with its own group
 			 * dependencies, add them as well. Not sure its
@@ -2411,14 +2412,8 @@ int netlink_nexthop_change(struct nlmsghdr *h, ns_id_t ns_id, int startup)
 			/* This is a change to a group we already have
 			 */
 
-			// TODO: Fix this for routes referencing it
-			/* Free what's already there */
-			zebra_nhg_free_members(nhe);
-
-			/* Update with new info */
-			nhe->nhg = nhg;
-			if (dep_count)
-				nhe->nhg_depends = nhg_depends;
+			zebra_nhg_set_invalid(nhe);
+			zebra_nhg_free_group_depends(&nhg, &nhg_depends);
 
 		} else {
 			/* This is a new nexthop group */
@@ -2435,6 +2430,7 @@ int netlink_nexthop_change(struct nlmsghdr *h, ns_id_t ns_id, int startup)
 			}
 
 			nhe->is_kernel_nh = true;
+
 			if (id != nhe->id) {
 				/* Duplicate but with different ID from
 				 * the kernel */
@@ -2449,18 +2445,18 @@ int netlink_nexthop_change(struct nlmsghdr *h, ns_id_t ns_id, int startup)
 					EC_ZEBRA_DUPLICATE_NHG_MESSAGE,
 					"Nexthop Group from kernel with ID (%d) is a duplicate, ignoring",
 					id);
-				zebra_nhg_connected_head_free(&nhg_depends);
+				nhg_connected_head_free(&nhg_depends);
+			} else {
+				/* Add the nhe to the interface's tree
+				 * of connected nhe's
+				 */
+				if (ifp)
+					zebra_nhg_set_if(nhe, ifp);
+
+				SET_FLAG(nhe->flags, NEXTHOP_GROUP_INSTALLED);
+				SET_FLAG(nhe->flags, NEXTHOP_GROUP_VALID);
 			}
 		}
-		if (ifp) {
-			/* Add the nhe to the interface's list
-			 * of connected nhe's
-			 */
-			// TODO: Don't add dupes
-			nhe_connected_add(ifp, nhe);
-		}
-		SET_FLAG(nhe->flags, NEXTHOP_GROUP_INSTALLED);
-		SET_FLAG(nhe->flags, NEXTHOP_GROUP_VALID);
 
 	} else if (h->nlmsg_type == RTM_DELNEXTHOP) {
 		if (!nhe) {
@@ -2471,9 +2467,10 @@ int netlink_nexthop_change(struct nlmsghdr *h, ns_id_t ns_id, int startup)
 			return -1;
 		}
 
-		UNSET_FLAG(nhe->flags, NEXTHOP_GROUP_INSTALLED);
+		zebra_nhg_set_invalid(nhe);
 
-		// TODO: Run some active check on all route_entry's?
+		// TODO: Probably won't need this if we expect
+		// upper level protocol to fix it.
 		if (nhe->refcnt) {
 			flog_err(
 				EC_ZEBRA_NHG_SYNC,
diff --git a/zebra/zebra_nhg.c b/zebra/zebra_nhg.c
index d3204bdf4..44265bc58 100644
--- a/zebra/zebra_nhg.c
+++ b/zebra/zebra_nhg.c
@@ -40,22 +40,22 @@
 #include "zebra/rt.h"
 #include "zebra_errors.h"
 #include "zebra_dplane.h"
+#include "zebra/interface.h"
 
 DEFINE_MTYPE_STATIC(ZEBRA, NHG, "Nexthop Group Entry");
 DEFINE_MTYPE_STATIC(ZEBRA, NHG_CONNECTED, "Nexthop Group Connected");
 
-static int zebra_nhg_connected_cmp(const struct nhg_connected *dep1,
-				   const struct nhg_connected *dep2);
+static int nhg_connected_cmp(const struct nhg_connected *dep1,
+			     const struct nhg_connected *dep2);
 
-RB_GENERATE(nhg_connected_head, nhg_connected, nhg_entry,
-	    zebra_nhg_connected_cmp);
+RB_GENERATE(nhg_connected_head, nhg_connected, nhg_entry, nhg_connected_cmp);
 
-static void nhg_connected_free(struct nhg_connected *dep)
+void nhg_connected_free(struct nhg_connected *dep)
 {
 	XFREE(MTYPE_NHG_CONNECTED, dep);
 }
 
-static struct nhg_connected *nhg_connected_new(struct nhg_hash_entry *nhe)
+struct nhg_connected *nhg_connected_new(struct nhg_hash_entry *nhe)
 {
 	struct nhg_connected *new = NULL;
 
@@ -65,35 +65,42 @@ static struct nhg_connected *nhg_connected_new(struct nhg_hash_entry *nhe)
 	return new;
 }
 
-static uint8_t zebra_nhg_connected_count(const struct nhg_connected_head *head)
+void nhg_connected_head_init(struct nhg_connected_head *head)
+{
+	RB_INIT(nhg_connected_head, head);
+}
+
+void nhg_connected_head_free(struct nhg_connected_head *head)
 {
 	struct nhg_connected *rb_node_dep = NULL;
-	uint8_t i = 0;
+	struct nhg_connected *tmp = NULL;
 
-	RB_FOREACH (rb_node_dep, nhg_connected_head, head) {
-		i++;
+	if (!nhg_connected_head_is_empty(head)) {
+		RB_FOREACH_SAFE (rb_node_dep, nhg_connected_head, head, tmp) {
+			RB_REMOVE(nhg_connected_head, head, rb_node_dep);
+			nhg_connected_free(rb_node_dep);
+		}
 	}
-	return i;
 }
 
-uint8_t zebra_nhg_depends_count(const struct nhg_hash_entry *nhe)
+unsigned int nhg_connected_head_count(const struct nhg_connected_head *head)
 {
-	return zebra_nhg_connected_count(&nhe->nhg_depends);
-}
+	struct nhg_connected *rb_node_dep = NULL;
+	unsigned int i = 0;
 
-static bool
-zebra_nhg_connected_head_is_empty(const struct nhg_connected_head *head)
-{
-	return RB_EMPTY(nhg_connected_head, head);
+	RB_FOREACH (rb_node_dep, nhg_connected_head, head) {
+		i++;
+	}
+	return i;
 }
 
-bool zebra_nhg_depends_is_empty(const struct nhg_hash_entry *nhe)
+bool nhg_connected_head_is_empty(const struct nhg_connected_head *head)
 {
-	return zebra_nhg_connected_head_is_empty(&nhe->nhg_depends);
+	return RB_EMPTY(nhg_connected_head, head);
 }
 
-void zebra_nhg_connected_head_del(struct nhg_connected_head *head,
-				  struct nhg_hash_entry *depend)
+void nhg_connected_head_del(struct nhg_connected_head *head,
+			    struct nhg_hash_entry *depend)
 {
 	struct nhg_connected lookup = {};
 	struct nhg_connected *removed = NULL;
@@ -105,8 +112,8 @@ void zebra_nhg_connected_head_del(struct nhg_connected_head *head,
 	nhg_connected_free(removed);
 }
 
-void zebra_nhg_connected_head_add(struct nhg_connected_head *head,
-				  struct nhg_hash_entry *depend)
+void nhg_connected_head_add(struct nhg_connected_head *head,
+			    struct nhg_hash_entry *depend)
 {
 	struct nhg_connected *new = NULL;
 
@@ -115,28 +122,14 @@ void zebra_nhg_connected_head_add(struct nhg_connected_head *head,
 	RB_INSERT(nhg_connected_head, head, new);
 }
 
-/**
- * zebra_nhg_dependents_del() - Delete a dependent from the nhg_hash_entry
- *
- * @from:	Nexthop group hash entry we are deleting from
- * @dependent:	Dependent we are deleting
- */
-void zebra_nhg_dependents_del(struct nhg_hash_entry *from,
-			      struct nhg_hash_entry *dependent)
+unsigned int zebra_nhg_depends_count(const struct nhg_hash_entry *nhe)
 {
-	zebra_nhg_connected_head_del(&from->nhg_dependents, dependent);
+	return nhg_connected_head_count(&nhe->nhg_depends);
 }
 
-/**
- * zebra_nhg_dependents_add() - Add a new dependent to the nhg_hash_entry
- *
- * @to:		Nexthop group hash entry we are adding to
- * @dependent:	Dependent we are adding
- */
-void zebra_nhg_dependents_add(struct nhg_hash_entry *to,
-			      struct nhg_hash_entry *dependent)
+bool zebra_nhg_depends_is_empty(const struct nhg_hash_entry *nhe)
 {
-	zebra_nhg_connected_head_add(&to->nhg_dependents, dependent);
+	return nhg_connected_head_is_empty(&nhe->nhg_depends);
 }
 
 /**
@@ -148,7 +141,7 @@ void zebra_nhg_dependents_add(struct nhg_hash_entry *to,
 void zebra_nhg_depends_del(struct nhg_hash_entry *from,
 			   struct nhg_hash_entry *depend)
 {
-	zebra_nhg_connected_head_del(&from->nhg_depends, depend);
+	nhg_connected_head_del(&from->nhg_depends, depend);
 
 	/* Delete from the dependent tree */
 	zebra_nhg_dependents_del(depend, from);
@@ -163,12 +156,7 @@ void zebra_nhg_depends_del(struct nhg_hash_entry *from,
 void zebra_nhg_depends_add(struct nhg_hash_entry *to,
 			   struct nhg_hash_entry *depend)
 {
-	zebra_nhg_connected_head_add(&to->nhg_depends, depend);
-}
-
-void zebra_nhg_connected_head_init(struct nhg_connected_head *head)
-{
-	RB_INIT(nhg_connected_head, head);
+	nhg_connected_head_add(&to->nhg_depends, depend);
 }
 
 /**
@@ -178,10 +166,9 @@ void zebra_nhg_connected_head_init(struct nhg_connected_head *head)
  */
 void zebra_nhg_depends_init(struct nhg_hash_entry *nhe)
 {
-	zebra_nhg_connected_head_init(&nhe->nhg_depends);
+	nhg_connected_head_init(&nhe->nhg_depends);
 }
 
-
 /**
  * zebra_nhg_depends_equal() - Are the dependencies of these nhe's equal
  *
@@ -217,6 +204,50 @@ static bool zebra_nhg_depends_equal(const struct nhg_hash_entry *nhe1,
 	return true;
 }
 
+unsigned int zebra_nhg_dependents_count(const struct nhg_hash_entry *nhe)
+{
+	return nhg_connected_head_count(&nhe->nhg_dependents);
+}
+
+bool zebra_nhg_dependents_is_empty(const struct nhg_hash_entry *nhe)
+{
+	return nhg_connected_head_is_empty(&nhe->nhg_dependents);
+}
+
+/**
+ * zebra_nhg_dependents_del() - Delete a dependent from the nhg_hash_entry
+ *
+ * @from:	Nexthop group hash entry we are deleting from
+ * @dependent:	Dependent we are deleting
+ */
+void zebra_nhg_dependents_del(struct nhg_hash_entry *from,
+			      struct nhg_hash_entry *dependent)
+{
+	nhg_connected_head_del(&from->nhg_dependents, dependent);
+}
+
+/**
+ * zebra_nhg_dependents_add() - Add a new dependent to the nhg_hash_entry
+ *
+ * @to:		Nexthop group hash entry we are adding to
+ * @dependent:	Dependent we are adding
+ */
+void zebra_nhg_dependents_add(struct nhg_hash_entry *to,
+			      struct nhg_hash_entry *dependent)
+{
+	nhg_connected_head_add(&to->nhg_dependents, dependent);
+}
+
+/**
+ * zebra_nhg_dependents_init() - Initialize tree for nhg dependents
+ *
+ * @nhe:	Nexthop group hash entry
+ */
+void zebra_nhg_dependents_init(struct nhg_hash_entry *nhe)
+{
+	nhg_connected_head_init(&nhe->nhg_dependents);
+}
+
 /**
  * zebra_nhg_lookup_id() - Lookup the nexthop group id in the id table
  *
@@ -280,6 +311,7 @@ static void *zebra_nhg_alloc(void *arg)
 	nhe->ifp = NULL;
 
 	/* Attach backpointer to anything that it depends on */
+	zebra_nhg_dependents_init(nhe);
 	if (!zebra_nhg_depends_is_empty(nhe)) {
 		RB_FOREACH (rb_node_dep, nhg_connected_head,
 			    &nhe->nhg_depends) {
@@ -291,6 +323,10 @@ static void *zebra_nhg_alloc(void *arg)
 	zebra_nhg_insert_id(nhe);
 
 
+	// TODO: This needs to be moved
+	// It should only install AFTER it gets
+	// the ifp right?
+	//
 	/* Send it to the kernel */
 	if (!nhe->is_kernel_nh)
 		zebra_nhg_install_kernel(nhe);
@@ -362,8 +398,8 @@ static int zebra_nhg_cmp(const struct nhg_hash_entry *nhe1,
 	return nhe1->id - nhe2->id;
 }
 
-static int zebra_nhg_connected_cmp(const struct nhg_connected *dep1,
-				   const struct nhg_connected *dep2)
+static int nhg_connected_cmp(const struct nhg_connected *dep1,
+			     const struct nhg_connected *dep2)
 {
 	return zebra_nhg_cmp(dep1->nhe, dep2->nhe);
 }
@@ -475,19 +511,6 @@ struct nhg_hash_entry *zebra_nhg_find_nexthop(struct nexthop *nh, afi_t afi)
 	return nhe;
 }
 
-void zebra_nhg_connected_head_free(struct nhg_connected_head *head)
-{
-	struct nhg_connected *rb_node_dep = NULL;
-	struct nhg_connected *tmp = NULL;
-
-	if (!zebra_nhg_connected_head_is_empty(head)) {
-		RB_FOREACH_SAFE (rb_node_dep, nhg_connected_head, head, tmp) {
-			RB_REMOVE(nhg_connected_head, head, rb_node_dep);
-			nhg_connected_free(rb_node_dep);
-		}
-	}
-}
-
 /**
  * zebra_nhg_free_group_depends() - Helper function for freeing nexthop_group
  * 				    struct and depends
@@ -508,7 +531,7 @@ void zebra_nhg_free_group_depends(struct nexthop_group **nhg,
 	//
 	//
 	if (head)
-		zebra_nhg_connected_head_free(head);
+		nhg_connected_head_free(head);
 
 	if (nhg)
 		nexthop_group_free_delete(nhg);
@@ -526,7 +549,7 @@ void zebra_nhg_free_members(struct nhg_hash_entry *nhe)
 	zebra_nhg_free_group_depends(&nhe->nhg, &nhe->nhg_depends);
 
 	// TODO: Fixup this function
-	zebra_nhg_connected_head_free(&nhe->nhg_dependents);
+	nhg_connected_head_free(&nhe->nhg_dependents);
 }
 
 /**
@@ -561,17 +584,6 @@ void zebra_nhg_release(struct nhg_hash_entry *nhe)
 }
 
 /**
- * zebra_nhg_uninstall_release() - Unistall and release a nhe
- *
- * @nhe:	Nexthop group hash entry
- */
-static void zebra_nhg_uninstall_release(struct nhg_hash_entry *nhe)
-{
-	zebra_nhg_uninstall_kernel(nhe);
-	// zebra_nhg_release(nhe);
-}
-
-/**
  * zebra_nhg_decrement_ref() - Decrement the reference count, release if unused
  *
  * @nhe:	Nexthop group hash entry
@@ -616,6 +628,29 @@ void zebra_nhg_increment_ref(struct nhg_hash_entry *nhe)
 	nhe->refcnt++;
 }
 
+void zebra_nhg_set_invalid(struct nhg_hash_entry *nhe)
+{
+
+	if (!zebra_nhg_dependents_is_empty(nhe)) {
+		struct nhg_connected *rb_node_dep = NULL;
+
+		RB_FOREACH (rb_node_dep, nhg_connected_head,
+			    &nhe->nhg_dependents) {
+			zebra_nhg_set_invalid(rb_node_dep->nhe);
+		}
+	}
+
+	UNSET_FLAG(nhe->flags, NEXTHOP_GROUP_VALID);
+	/* Assuming uninstalled as well here */
+	UNSET_FLAG(nhe->flags, NEXTHOP_GROUP_INSTALLED);
+}
+
+void zebra_nhg_set_if(struct nhg_hash_entry *nhe, struct interface *ifp)
+{
+	nhe->ifp = ifp;
+	if_nhg_dependents_add(ifp, nhe);
+}
+
 static void nexthop_set_resolved(afi_t afi, const struct nexthop *newhop,
 				 struct nexthop *nexthop)
 {
diff --git a/zebra/zebra_nhg.h b/zebra/zebra_nhg.h
index 4aa59e753..ebdc3a636 100644
--- a/zebra/zebra_nhg.h
+++ b/zebra/zebra_nhg.h
@@ -96,30 +96,46 @@ struct nhg_connected {
 	struct nhg_hash_entry *nhe;
 };
 
-RB_PROTOTYPE(nhg_connected_head, nhg_connected, nhg_entry,
-	     zebra_nhg_connected_cmp);
+RB_PROTOTYPE(nhg_connected_head, nhg_connected, nhg_entry, nhg_connected_cmp);
 
 void zebra_nhg_init(void);
 void zebra_nhg_terminate(void);
 
-extern uint8_t zebra_nhg_depends_count(const struct nhg_hash_entry *nhe);
+extern void nhg_connected_free(struct nhg_connected *dep);
+extern struct nhg_connected *nhg_connected_new(struct nhg_hash_entry *nhe);
+
+/* nhg connected tree direct access functions */
+extern void nhg_connected_head_init(struct nhg_connected_head *head);
+extern unsigned int
+nhg_connected_head_count(const struct nhg_connected_head *head);
+extern void nhg_connected_head_free(struct nhg_connected_head *head);
+extern bool nhg_connected_head_is_empty(const struct nhg_connected_head *head);
+extern void nhg_connected_head_del(struct nhg_connected_head *head,
+				   struct nhg_hash_entry *nhe);
+extern void nhg_connected_head_add(struct nhg_connected_head *head,
+				   struct nhg_hash_entry *nhe);
+
+/**
+ * NHE abstracted tree functions.
+ * Use these where possible instead of the direct ones access ones.
+ */
+/* Depends */
+extern unsigned int zebra_nhg_depends_count(const struct nhg_hash_entry *nhe);
 extern bool zebra_nhg_depends_is_empty(const struct nhg_hash_entry *nhe);
-extern void zebra_nhg_connected_head_del(struct nhg_connected_head *head,
-					 struct nhg_hash_entry *nhe);
-extern void zebra_nhg_connected_head_add(struct nhg_connected_head *head,
-					 struct nhg_hash_entry *nhe);
-extern void zebra_nhg_dependents_del(struct nhg_hash_entry *from,
-				     struct nhg_hash_entry *dependent);
-extern void zebra_nhg_dependents_add(struct nhg_hash_entry *to,
-				     struct nhg_hash_entry *dependent);
 extern void zebra_nhg_depends_del(struct nhg_hash_entry *from,
 				  struct nhg_hash_entry *depend);
 extern void zebra_nhg_depends_add(struct nhg_hash_entry *to,
 				  struct nhg_hash_entry *depend);
-extern void zebra_nhg_connected_head_init(struct nhg_connected_head *head);
 extern void zebra_nhg_depends_init(struct nhg_hash_entry *nhe);
-extern void zebra_nhg_depends_copy(struct nhg_hash_entry *to,
-				   struct nhg_hash_entry *from);
+/* Dependents */
+extern unsigned int
+zebra_nhg_dependents_count(const struct nhg_hash_entry *nhe);
+extern bool zebra_nhg_dependents_is_empty(const struct nhg_hash_entry *nhe);
+extern void zebra_nhg_dependents_del(struct nhg_hash_entry *from,
+				     struct nhg_hash_entry *dependent);
+extern void zebra_nhg_dependents_add(struct nhg_hash_entry *to,
+				     struct nhg_hash_entry *dependent);
+extern void zebra_nhg_dependents_init(struct nhg_hash_entry *nhe);
 
 
 extern struct nhg_hash_entry *zebra_nhg_lookup_id(uint32_t id);
@@ -140,7 +156,6 @@ extern struct nhg_hash_entry *zebra_nhg_find_nexthop(struct nexthop *nh,
 						     afi_t afi);
 
 
-void zebra_nhg_connected_head_free(struct nhg_connected_head *head);
 void zebra_nhg_free_group_depends(struct nexthop_group **nhg,
 				  struct nhg_connected_head *head);
 void zebra_nhg_free_members(struct nhg_hash_entry *nhe);
@@ -148,6 +163,8 @@ void zebra_nhg_free(void *arg);
 void zebra_nhg_release(struct nhg_hash_entry *nhe);
 void zebra_nhg_decrement_ref(struct nhg_hash_entry *nhe);
 void zebra_nhg_increment_ref(struct nhg_hash_entry *nhe);
+void zebra_nhg_set_invalid(struct nhg_hash_entry *nhe);
+void zebra_nhg_set_if(struct nhg_hash_entry *nhe, struct interface *ifp);
 
 extern int nexthop_active_update(struct route_node *rn, struct route_entry *re);
 
diff --git a/zebra/zebra_rib.c b/zebra/zebra_rib.c
index 7015ff77d..898076b01 100644
--- a/zebra/zebra_rib.c
+++ b/zebra/zebra_rib.c
@@ -2667,7 +2667,7 @@ int rib_add_multipath(afi_t afi, safi_t safi, struct prefix *p,
 		struct nexthop lookup = {0};
 		struct nhg_hash_entry *depend = NULL;
 
-		zebra_nhg_connected_head_init(&nhg_depends);
+		nhg_connected_head_init(&nhg_depends);
 
 		for (ALL_NEXTHOPS_PTR(re->ng, nh)) {
 			lookup = *nh;
@@ -2675,7 +2675,7 @@ int rib_add_multipath(afi_t afi, safi_t safi, struct prefix *p,
 			lookup.next = NULL;
 			/* Use the route afi here, since a single nh */
 			depend = zebra_nhg_find_nexthop(&lookup, afi);
-			zebra_nhg_connected_head_add(&nhg_depends, depend);
+			nhg_connected_head_add(&nhg_depends, depend);
 		}
 
 		/* change the afi for group */
@@ -2698,7 +2698,7 @@ int rib_add_multipath(afi_t afi, safi_t safi, struct prefix *p,
 			EC_ZEBRA_TABLE_LOOKUP_FAILED,
 			"Zebra failed to find or create a nexthop hash entry for id=%u in a route entry",
 			re->nhe_id);
-		zebra_nhg_connected_head_free(&nhg_depends);
+		nhg_connected_head_free(&nhg_depends);
 	}