From 0c8215cbab59480616fe31d8a81323070a21137b Mon Sep 17 00:00:00 2001
From: Stephen Worley <sworley@cumulusnetworks.com>
Date: Mon, 13 May 2019 17:10:34 -0700
Subject: [PATCH] zebra,lib: Refactor depends to RB tree

Refactor the depends to use an RB tree instead of a list.

Signed-off-by: Stephen Worley <sworley@cumulusnetworks.com>
---
 lib/nexthop_group.c  |   7 +
 lib/nexthop_group.h  |   1 +
 zebra/rt_netlink.c   |  70 +++++-----
 zebra/zebra_dplane.c | 105 +++++++++++----
 zebra/zebra_dplane.h |  11 +-
 zebra/zebra_nhg.c    | 308 +++++++++++++++++++++++--------------------
 zebra/zebra_nhg.h    |  54 ++++++--
 zebra/zebra_rib.c    |  20 +--
 zebra/zebra_vty.c    |  12 +-
 9 files changed, 359 insertions(+), 229 deletions(-)

diff --git a/lib/nexthop_group.c b/lib/nexthop_group.c
index 9564321d38..527039ac63 100644
--- a/lib/nexthop_group.c
+++ b/lib/nexthop_group.c
@@ -122,6 +122,13 @@ void nexthop_group_delete(struct nexthop_group **nhg)
 	XFREE(MTYPE_NEXTHOP_GROUP, *nhg);
 }
 
+void nexthop_group_free_delete(struct nexthop_group **nhg)
+{
+	if ((*nhg)->nexthop)
+		nexthops_free((*nhg)->nexthop);
+	nexthop_group_delete(nhg);
+}
+
 /* Add nexthop to the end of a nexthop list.  */
 void _nexthop_add(struct nexthop **target, struct nexthop *nexthop)
 {
diff --git a/lib/nexthop_group.h b/lib/nexthop_group.h
index 4f4d40eb33..5193c943c4 100644
--- a/lib/nexthop_group.h
+++ b/lib/nexthop_group.h
@@ -41,6 +41,7 @@ struct nexthop_group {
 
 struct nexthop_group *nexthop_group_new(void);
 void nexthop_group_delete(struct nexthop_group **nhg);
+void nexthop_group_free_delete(struct nexthop_group **nhg);
 
 void nexthop_group_copy(struct nexthop_group *to,
 			struct nexthop_group *from);
diff --git a/zebra/rt_netlink.c b/zebra/rt_netlink.c
index 7917e3046e..f7e72e243f 100644
--- a/zebra/rt_netlink.c
+++ b/zebra/rt_netlink.c
@@ -1888,26 +1888,24 @@ int kernel_get_ipmr_sg_stats(struct zebra_vrf *zvrf, void *in)
  *
  * @n:			Netlink message header struct
  * @req_size:		Size allocated for this message
- * @nhg_depends:	List of entry dependencies
+ * @depends_info:	Array of depend_info structs
+ * @count:		How many depencies there are
  */
 static void _netlink_nexthop_build_group(struct nlmsghdr *n, size_t req_size,
-					 struct list *nhg_depends)
+					 const struct depend_info *depends_info,
+					 const uint8_t count)
 {
-	int count = listcount(nhg_depends);
-	int i = 0;
 	struct nexthop_grp grp[count];
-	struct listnode *dp_node = NULL;
-	struct nhg_depend *n_dp = NULL;
 
 	memset(grp, 0, sizeof(grp));
 
 	if (count) {
-		for (ALL_LIST_ELEMENTS_RO(nhg_depends, dp_node, n_dp)) {
-			grp[i++].id = n_dp->nhe->id;
+		for (int i = 0; i < count; i++) {
+			grp[i].id = depends_info[i].id;
+			grp[i].weight = depends_info[i].weight;
 		}
+		addattr_l(n, req_size, NHA_GROUP, grp, count * sizeof(*grp));
 	}
-
-	addattr_l(n, req_size, NHA_GROUP, grp, count * sizeof(*grp));
 }
 
 /**
@@ -1927,8 +1925,6 @@ static int netlink_nexthop(int cmd, struct zebra_dplane_ctx *ctx)
 
 	memset(&req, 0, sizeof(req));
 
-	const struct nhg_hash_entry *nhe = dplane_ctx_get_nhe(ctx);
-
 	req.n.nlmsg_len = NLMSG_LENGTH(sizeof(struct nhmsg));
 	req.n.nlmsg_flags = NLM_F_CREATE | NLM_F_REQUEST;
 	req.n.nlmsg_type = cmd;
@@ -1937,25 +1933,31 @@ static int netlink_nexthop(int cmd, struct zebra_dplane_ctx *ctx)
 	req.nhm.nh_family = AF_UNSPEC;
 	// TODO: Scope?
 
-	if (!nhe->id) {
+	uint32_t id = dplane_ctx_get_nhe_id(ctx);
+
+	if (!id) {
 		flog_err(
 			EC_ZEBRA_NHG_FIB_UPDATE,
 			"Failed trying to update a nexthop group in the kernel that does not have an ID");
 		return -1;
 	}
 
-	addattr32(&req.n, sizeof(req), NHA_ID, nhe->id);
+	addattr32(&req.n, sizeof(req), NHA_ID, id);
 
 	if (cmd == RTM_NEWNEXTHOP) {
-		if (nhe->nhg_depends) {
-			_netlink_nexthop_build_group(&req.n, sizeof(req),
-						     nhe->nhg_depends);
-		} else {
-			const struct nexthop *nh = nhe->nhg->nexthop;
+		if (dplane_ctx_get_nhe_depends_count(ctx))
+			_netlink_nexthop_build_group(
+				&req.n, sizeof(req),
+				dplane_ctx_get_nhe_depends_info(ctx),
+				dplane_ctx_get_nhe_depends_count(ctx));
+		else {
+			const struct nexthop *nh =
+				dplane_ctx_get_nhe_ng(ctx)->nexthop;
+			afi_t afi = dplane_ctx_get_nhe_afi(ctx);
 
-			if (nhe->afi == AFI_IP)
+			if (afi == AFI_IP)
 				req.nhm.nh_family = AF_INET;
-			else if (nhe->afi == AFI_IP6)
+			else if (afi == AFI_IP6)
 				req.nhm.nh_family = AF_INET6;
 
 			switch (nh->type) {
@@ -1998,7 +2000,7 @@ static int netlink_nexthop(int cmd, struct zebra_dplane_ctx *ctx)
 		return -1;
 	}
 
-	_netlink_nexthop_debug(cmd, nhe->id);
+	_netlink_nexthop_debug(cmd, id);
 
 	return netlink_talk_info(netlink_talk_filter, &req.n,
 				 dplane_ctx_get_ns(ctx), 0);
@@ -2223,14 +2225,14 @@ static struct nexthop *netlink_nexthop_process_nh(struct rtattr **tb,
  * netlink_nexthop_process_group() - Iterate over nhmsg nexthop group
  *
  * @tb:			Netlink RTA data
- * @nhg_depends:	List of nexthops in the group (depends)
+ * @nhg_depends:	Tree head of nexthops in the group (depends)
  * @nhg:		Nexthop group struct
  *
  * Return:	Count of nexthops in the group
  */
 static int netlink_nexthop_process_group(struct rtattr **tb,
 					 struct nexthop_group *nhg,
-					 struct list **nhg_depends)
+					 struct nhg_depends_head *nhg_depends)
 {
 	int count = 0;
 	struct nexthop_grp *n_grp = NULL;
@@ -2249,7 +2251,7 @@ static int netlink_nexthop_process_group(struct rtattr **tb,
 	zlog_debug("Nexthop group type: %d",
 		   *((uint16_t *)RTA_DATA(tb[NHA_GROUP_TYPE])));
 
-	*nhg_depends = nhg_depend_new_list();
+	zebra_nhg_depends_head_init(nhg_depends);
 
 	for (int i = 0; i < count; i++) {
 		/* We do not care about nexthop_grp.weight at
@@ -2259,7 +2261,7 @@ static int netlink_nexthop_process_group(struct rtattr **tb,
 		 */
 		depend = zebra_nhg_lookup_id(n_grp[i].id);
 		if (depend) {
-			nhg_depend_add(*nhg_depends, depend);
+			zebra_nhg_depends_head_add(nhg_depends, depend);
 			/*
 			 * If this is a nexthop with its own group
 			 * dependencies, add them as well. Not sure its
@@ -2301,8 +2303,8 @@ int netlink_nexthop_change(struct nlmsghdr *h, ns_id_t ns_id, int startup)
 	/* struct for nexthop group abstraction  */
 	struct nexthop_group *nhg = NULL;
 	struct nexthop *nh = NULL;
-	/* If its a group, array of nexthops */
-	struct list *nhg_depends = NULL;
+	/* If its a group, tree head of nexthops */
+	struct nhg_depends_head nhg_depends = {0};
 	/* Count of nexthops in group array */
 	int dep_count = 0;
 	/* struct that goes into our tables */
@@ -2396,11 +2398,16 @@ int netlink_nexthop_change(struct nlmsghdr *h, ns_id_t ns_id, int startup)
 
 		if (!nhg->nexthop) {
 			/* Nothing to lookup */
-			zebra_nhg_free_group_depends(nhg, nhg_depends);
+			zebra_nhg_free_group_depends(&nhg, &nhg_depends);
 			return -1;
 		}
 
 		if (nhe) {
+			// TODO: Apparently we don't want changes
+			// to already created one in our table.
+			// They should be immutable...
+			// Gotta figure that one out.
+
 			/* This is a change to a group we already have
 			 */
 
@@ -2415,9 +2422,9 @@ int netlink_nexthop_change(struct nlmsghdr *h, ns_id_t ns_id, int startup)
 
 		} else {
 			/* This is a new nexthop group */
-			nhe = zebra_nhg_find(nhg, vrf_id, afi, id, nhg_depends,
+			nhe = zebra_nhg_find(nhg, vrf_id, afi, id, &nhg_depends,
 					     true);
-			zebra_nhg_free_group_depends(nhg, nhg_depends);
+			nexthop_group_free_delete(&nhg);
 
 			if (!nhe) {
 				flog_err(
@@ -2477,7 +2484,6 @@ int netlink_nexthop_change(struct nlmsghdr *h, ns_id_t ns_id, int startup)
 		}
 	}
 
-
 	return 0;
 }
 
diff --git a/zebra/zebra_dplane.c b/zebra/zebra_dplane.c
index 57faca8bb4..170f7cfe64 100644
--- a/zebra/zebra_dplane.c
+++ b/zebra/zebra_dplane.c
@@ -66,6 +66,20 @@ const uint32_t DPLANE_DEFAULT_NEW_WORK = 100;
 
 #endif	/* DPLANE_DEBUG */
 
+/*
+ * Nexthop information captured for nexthop/nexthop group updates
+ */
+struct dplane_nexthop_info {
+	uint32_t id;
+	afi_t afi;
+	vrf_id_t vrf_id;
+	bool is_kernel_nh;
+
+	struct nexthop_group ng;
+	struct depend_info depends_info[MULTIPATH_NUM];
+	uint8_t depends_count;
+};
+
 /*
  * Route information captured for route updates.
  */
@@ -95,8 +109,8 @@ struct dplane_route_info {
 	uint32_t zd_mtu;
 	uint32_t zd_nexthop_mtu;
 
-	/* Nexthop hash entry */
-	struct nhg_hash_entry zd_nhe;
+	/* Nexthop hash entry info */
+	struct dplane_nexthop_info nhe;
 
 	/* Nexthops */
 	struct nexthop_group zd_ng;
@@ -470,7 +484,12 @@ static void dplane_ctx_free(struct zebra_dplane_ctx **pctx)
 	case DPLANE_OP_NH_INSTALL:
 	case DPLANE_OP_NH_UPDATE:
 	case DPLANE_OP_NH_DELETE: {
-		zebra_nhg_free_members(&(*pctx)->u.rinfo.zd_nhe);
+		if ((*pctx)->u.rinfo.nhe.ng.nexthop) {
+			/* This deals with recursive nexthops too */
+			nexthops_free((*pctx)->u.rinfo.nhe.ng.nexthop);
+
+			(*pctx)->u.rinfo.nhe.ng.nexthop = NULL;
+		}
 		break;
 	}
 
@@ -1040,11 +1059,48 @@ const struct zebra_dplane_info *dplane_ctx_get_ns(
 }
 
 /* Accessors for nexthop information */
-const struct nhg_hash_entry *
-dplane_ctx_get_nhe(const struct zebra_dplane_ctx *ctx)
+uint32_t dplane_ctx_get_nhe_id(const struct zebra_dplane_ctx *ctx)
+{
+	DPLANE_CTX_VALID(ctx);
+	return ctx->u.rinfo.nhe.id;
+}
+
+afi_t dplane_ctx_get_nhe_afi(const struct zebra_dplane_ctx *ctx)
+{
+	DPLANE_CTX_VALID(ctx);
+	return ctx->u.rinfo.nhe.afi;
+}
+
+vrf_id_t dplane_ctx_get_nhe_vrf_id(const struct zebra_dplane_ctx *ctx)
 {
 	DPLANE_CTX_VALID(ctx);
-	return &(ctx->u.rinfo.zd_nhe);
+	return ctx->u.rinfo.nhe.vrf_id;
+}
+
+bool dplane_ctx_get_nhe_is_kernel_nh(const struct zebra_dplane_ctx *ctx)
+{
+	DPLANE_CTX_VALID(ctx);
+	return ctx->u.rinfo.nhe.is_kernel_nh;
+}
+
+const struct nexthop_group *
+dplane_ctx_get_nhe_ng(const struct zebra_dplane_ctx *ctx)
+{
+	DPLANE_CTX_VALID(ctx);
+	return &(ctx->u.rinfo.nhe.ng);
+}
+
+const struct depend_info *
+dplane_ctx_get_nhe_depends_info(const struct zebra_dplane_ctx *ctx)
+{
+	DPLANE_CTX_VALID(ctx);
+	return ctx->u.rinfo.nhe.depends_info;
+}
+
+uint8_t dplane_ctx_get_nhe_depends_count(const struct zebra_dplane_ctx *ctx)
+{
+	DPLANE_CTX_VALID(ctx);
+	return ctx->u.rinfo.nhe.depends_count;
 }
 
 /* Accessors for LSP information */
@@ -1505,21 +1561,26 @@ static int dplane_ctx_nexthop_init(struct zebra_dplane_ctx *ctx,
 	ctx->zd_status = ZEBRA_DPLANE_REQUEST_SUCCESS;
 
 	/* Copy over nhe info */
-	ctx->u.rinfo.zd_nhe.id = nhe->id;
-	ctx->u.rinfo.zd_nhe.vrf_id = nhe->vrf_id;
-	ctx->u.rinfo.zd_nhe.afi = nhe->afi;
-	ctx->u.rinfo.zd_nhe.refcnt = nhe->refcnt;
-	ctx->u.rinfo.zd_nhe.is_kernel_nh = nhe->is_kernel_nh;
-	ctx->u.rinfo.zd_nhe.dplane_ref = nhe->dplane_ref;
-	ctx->u.rinfo.zd_nhe.ifp = nhe->ifp;
-
-	ctx->u.rinfo.zd_nhe.nhg = nexthop_group_new();
-	nexthop_group_copy(ctx->u.rinfo.zd_nhe.nhg, nhe->nhg);
-
-	if (nhe->nhg_depends)
-		ctx->u.rinfo.zd_nhe.nhg_depends =
-			nhg_depend_dup_list(nhe->nhg_depends);
-
+	ctx->u.rinfo.nhe.id = nhe->id;
+	ctx->u.rinfo.nhe.afi = nhe->afi;
+	ctx->u.rinfo.nhe.vrf_id = nhe->vrf_id;
+	ctx->u.rinfo.nhe.is_kernel_nh = nhe->is_kernel_nh;
+
+	nexthop_group_copy(&(ctx->u.rinfo.nhe.ng), nhe->nhg);
+
+	if (!zebra_nhg_depends_is_empty(nhe)) {
+		struct nhg_depend *rb_node_dep = NULL;
+		uint8_t i = 0;
+
+		RB_FOREACH (rb_node_dep, nhg_depends_head, &nhe->nhg_depends) {
+			ctx->u.rinfo.nhe.depends_info[i].id =
+				rb_node_dep->nhe->id;
+			/* We aren't using weights for anything right now */
+			ctx->u.rinfo.nhe.depends_info[i].weight = 0;
+			i++;
+		}
+		ctx->u.rinfo.nhe.depends_count = i;
+	}
 
 	/* Extract ns info - can't use pointers to 'core'
 	   structs */
@@ -3053,7 +3114,7 @@ kernel_dplane_nexthop_update(struct zebra_dplane_ctx *ctx)
 
 	if (IS_ZEBRA_DEBUG_DPLANE_DETAIL) {
 		zlog_debug("ID (%u) Dplane nexthop update ctx %p op %s",
-			   dplane_ctx_get_nhe(ctx)->id, ctx,
+			   dplane_ctx_get_nhe_id(ctx), ctx,
 			   dplane_op2str(dplane_ctx_get_op(ctx)));
 	}
 
diff --git a/zebra/zebra_dplane.h b/zebra/zebra_dplane.h
index 6ca2a274c6..447e9c0d7f 100644
--- a/zebra/zebra_dplane.h
+++ b/zebra/zebra_dplane.h
@@ -276,8 +276,15 @@ const struct nexthop_group *dplane_ctx_get_old_ng(
 	const struct zebra_dplane_ctx *ctx);
 
 /* Accessors for nexthop information */
-const struct nhg_hash_entry *
-dplane_ctx_get_nhe(const struct zebra_dplane_ctx *ctx);
+uint32_t dplane_ctx_get_nhe_id(const struct zebra_dplane_ctx *ctx);
+afi_t dplane_ctx_get_nhe_afi(const struct zebra_dplane_ctx *ctx);
+vrf_id_t dplane_ctx_get_nhe_vrf_id(const struct zebra_dplane_ctx *ctx);
+bool dplane_ctx_get_nhe_is_kernel_nh(const struct zebra_dplane_ctx *ctx);
+const struct nexthop_group *
+dplane_ctx_get_nhe_ng(const struct zebra_dplane_ctx *ctx);
+const struct depend_info *
+dplane_ctx_get_nhe_depends_info(const struct zebra_dplane_ctx *ctx);
+uint8_t dplane_ctx_get_nhe_depends_count(const struct zebra_dplane_ctx *ctx);
 
 /* Accessors for LSP information */
 mpls_label_t dplane_ctx_get_in_label(const struct zebra_dplane_ctx *ctx);
diff --git a/zebra/zebra_nhg.c b/zebra/zebra_nhg.c
index 94565f434d..56f8c9d852 100644
--- a/zebra/zebra_nhg.c
+++ b/zebra/zebra_nhg.c
@@ -39,111 +39,133 @@
 #include "zebra/zserv.h"
 #include "zebra/rt.h"
 #include "zebra_errors.h"
+#include "zebra_dplane.h"
 
 DEFINE_MTYPE_STATIC(ZEBRA, NHG, "Nexthop Group Entry");
-DEFINE_MTYPE_STATIC(ZEBRA, NHG_DEPENDS, "Nexthop Group Entry Depends");
+DEFINE_MTYPE_STATIC(ZEBRA, NHG_DEPEND, "Nexthop Group Dependency");
+
+static int zebra_nhg_depend_cmp(const struct nhg_depend *dep1,
+				const struct nhg_depend *dep2);
+
+RB_GENERATE(nhg_depends_head, nhg_depend, depend, zebra_nhg_depend_cmp);
+
+static void nhg_depend_free(struct nhg_depend *dep)
+{
+	XFREE(MTYPE_NHG_DEPEND, dep);
+}
+
+static struct nhg_depend *nhg_depend_new(struct nhg_hash_entry *nhe)
+{
+	struct nhg_depend *new = NULL;
+
+	new = XCALLOC(MTYPE_NHG_DEPEND, sizeof(struct nhg_depend));
+	new->nhe = nhe;
+
+	return new;
+}
+
+static uint8_t zebra_nhg_depends_head_count(const struct nhg_depends_head *head)
+{
+	struct nhg_depend *rb_node_dep = NULL;
+	uint8_t i = 0;
+
+	RB_FOREACH (rb_node_dep, nhg_depends_head, head) {
+		i++;
+	}
+	return i;
+}
+
+uint8_t zebra_nhg_depends_count(const struct nhg_hash_entry *nhe)
+{
+	return zebra_nhg_depends_head_count(&nhe->nhg_depends);
+}
+
+static bool zebra_nhg_depends_head_is_empty(const struct nhg_depends_head *head)
+{
+	return RB_EMPTY(nhg_depends_head, head);
+}
 
 /**
- * nhg_depend_add() - Add a new dependency to the nhg_hash_entry
+ * zebra_nhg_depends_is_empty() - Are there any dependencies
  *
- * @nhg_depends:	List we are adding the dependency to
- * @depend:		Dependency we are adding
+ * @nhe:	Nexthop group hash entry
  *
- * Return:		Newly created nhg_depend
+ * Return:	True if empty, False otherwise
  */
-struct nhg_depend *nhg_depend_add(struct list *nhg_depends,
-				  struct nhg_hash_entry *depend)
+bool zebra_nhg_depends_is_empty(const struct nhg_hash_entry *nhe)
 {
-	struct nhg_depend *nhg_dp = NULL;
+	return zebra_nhg_depends_head_is_empty(&nhe->nhg_depends);
+}
+
+void zebra_nhg_depends_head_del(struct nhg_depends_head *head,
+			  struct nhg_hash_entry *depend)
+{
+	struct nhg_depend lookup = {};
+	struct nhg_depend *removed = NULL;
+
+	lookup.nhe = depend;
 
-	nhg_dp = nhg_depend_new();
-	nhg_dp->nhe = depend;
+	removed = RB_REMOVE(nhg_depends_head, head, &lookup);
 
-	listnode_add(nhg_depends, nhg_dp);
-	return nhg_dp;
+	nhg_depend_free(removed);
 }
 
-/**
- * nhg_depend_new() - Allocate a new nhg_depend struct
- *
- * Return:	Allocated nhg_depend struct
- */
-struct nhg_depend *nhg_depend_new(void)
+// TODO: Refactor this for use with dependents as well
+void zebra_nhg_depends_head_add(struct nhg_depends_head *head,
+			  struct nhg_hash_entry *depend)
 {
-	return XCALLOC(MTYPE_NHG_DEPENDS, sizeof(struct nhg_depend));
+	struct nhg_depend *new = NULL;
+
+	new = nhg_depend_new(depend);
+
+	RB_INSERT(nhg_depends_head, head, new);
 }
 
 /**
- * nhg_depend_free() - Free the nhg_depend struct
+ * zebra_nhg_depend_del() - Delete a dependency from the nhg_hash_entry
+ *
+ * @from:	Nexthop group hash entry we are deleting from
+ * @depend:	Dependency we are deleting
  */
-void nhg_depend_free(struct nhg_depend *depend)
+void zebra_nhg_depends_del(struct nhg_hash_entry *from,
+			   struct nhg_hash_entry *depend)
 {
-	XFREE(MTYPE_NHG_DEPENDS, depend);
+	zebra_nhg_depends_head_del(&from->nhg_depends, depend);
 }
 
 /**
- * nhg_depend_new_list() - Allocate a new list for nhg_depends
+ * zebra_nhg_depends_add() - Add a new dependency to the nhg_hash_entry
  *
- * Return:	Allocated nhg_depend list
+ * @to:		Nexthop group hash entry we are adding to
+ * @depend:	Dependency we are adding
  */
-struct list *nhg_depend_new_list()
+void zebra_nhg_depends_add(struct nhg_hash_entry *to,
+			   struct nhg_hash_entry *depend)
 {
-	struct list *nhg_depends = NULL;
-
-	nhg_depends = list_new();
-	nhg_depends->del = (void (*)(void *))nhg_depend_free;
-
-	return nhg_depends;
+	zebra_nhg_depends_head_add(&to->nhg_depends, depend);
 }
 
 /**
- * nhg_depend_dup_list() - Duplicate the dependency linked list
+ * zebra_nhg_depends_head_init() - Helper to init the RB tree head directly
  *
- * @from:	List to duplicate
- *
- * Return:	New list
+ * @head:	RB nhg_depends_head struct
  */
-struct list *nhg_depend_dup_list(struct list *from)
+void zebra_nhg_depends_head_init(struct nhg_depends_head *head)
 {
-	struct list *to = NULL;
-	struct listnode *ln = NULL;
-	struct nhg_depend *n_dp = NULL;
-
-	to = nhg_depend_new_list();
-
-	for (ALL_LIST_ELEMENTS_RO(from, ln, n_dp)) {
-		nhg_depend_add(to, n_dp->nhe);
-	}
-
-	return to;
+	RB_INIT(nhg_depends_head, head);
 }
 
 /**
- * zebra_nhg_depends_lookup_id() - Lookup for an id in the nhg_depends
- * 				   linked list of another nhe
+ * zebra_nhg_depends_init() - Initialize tree for nhg dependencies
  *
- * @nhe: 	Nexthop group hash entry with list to look in
- * @lookup:	ID to look for
- *
- * Return:	Nexthop group hash entry if found
+ * @nhe:	Nexthop group hash entry
  */
-static struct nhg_hash_entry *
-zebra_nhg_depends_lookup_id(const struct nhg_hash_entry *nhe, const uint32_t id)
+void zebra_nhg_depends_init(struct nhg_hash_entry *nhe)
 {
-	struct listnode *ln = NULL;
-	struct nhg_depend *n_dp = NULL;
-	struct nhg_hash_entry *match = NULL;
-
-	for (ALL_LIST_ELEMENTS_RO(nhe->nhg_depends, ln, n_dp)) {
-		if (n_dp->nhe->id == id) {
-			match = n_dp->nhe;
-			break;
-		}
-	}
-
-	return match;
+	zebra_nhg_depends_head_init(&nhe->nhg_depends);
 }
 
+
 /**
  * zebra_nhg_depends_equal() - Are the dependencies of these nhe's equal
  *
@@ -158,21 +180,20 @@ zebra_nhg_depends_lookup_id(const struct nhg_hash_entry *nhe, const uint32_t id)
 static bool zebra_nhg_depends_equal(const struct nhg_hash_entry *nhe1,
 				    const struct nhg_hash_entry *nhe2)
 {
-	struct listnode *ln = NULL;
-	struct nhg_depend *n_dp = NULL;
+	struct nhg_depend *rb_node_dep = NULL;
 
-	if (!nhe1->nhg_depends && !nhe2->nhg_depends)
+	if (zebra_nhg_depends_is_empty(nhe1)
+	    && zebra_nhg_depends_is_empty(nhe2))
 		return true;
 
-	if ((nhe1->nhg_depends && !nhe2->nhg_depends)
-	    || (nhe2->nhg_depends && !nhe1->nhg_depends))
-		return false;
-
-	if (listcount(nhe1->nhg_depends) != listcount(nhe2->nhg_depends))
+	if ((zebra_nhg_depends_is_empty(nhe1)
+	     && !zebra_nhg_depends_is_empty(nhe2))
+	    || (zebra_nhg_depends_is_empty(nhe2)
+		&& !zebra_nhg_depends_is_empty(nhe1)))
 		return false;
 
-	for (ALL_LIST_ELEMENTS_RO(nhe1->nhg_depends, ln, n_dp)) {
-		if (!zebra_nhg_depends_lookup_id(nhe2, n_dp->nhe->id))
+	RB_FOREACH (rb_node_dep, nhg_depends_head, &nhe1->nhg_depends) {
+		if (!RB_FIND(nhg_depends_head, &nhe2->nhg_depends, rb_node_dep))
 			return false;
 	}
 
@@ -188,7 +209,7 @@ static bool zebra_nhg_depends_equal(const struct nhg_hash_entry *nhe1,
  */
 struct nhg_hash_entry *zebra_nhg_lookup_id(uint32_t id)
 {
-	struct nhg_hash_entry lookup = {0};
+	struct nhg_hash_entry lookup = {};
 
 	lookup.id = id;
 	return hash_lookup(zrouter.nhgs_id, &lookup);
@@ -227,10 +248,7 @@ static void *zebra_nhg_alloc(void *arg)
 
 	nhe->id = copy->id;
 
-	nhe->nhg_depends = NULL;
-
-	if (copy->nhg_depends)
-		nhe->nhg_depends = nhg_depend_dup_list(copy->nhg_depends);
+	nhe->nhg_depends = copy->nhg_depends;
 
 	nhe->nhg = nexthop_group_new();
 	nexthop_group_copy(nhe->nhg, copy->nhg);
@@ -254,42 +272,6 @@ static void *zebra_nhg_alloc(void *arg)
 	return nhe;
 }
 
-static uint32_t zebra_nhg_hash_key_nexthop_group(struct nexthop_group *nhg)
-{
-	struct nexthop *nh;
-	uint32_t i;
-	uint32_t key = 0;
-
-	/*
-	 * We are not interested in hashing over any recursively
-	 * resolved nexthops
-	 */
-	for (nh = nhg->nexthop; nh; nh = nh->next) {
-		key = jhash_1word(nh->type, key);
-		key = jhash_2words(nh->vrf_id, nh->nh_label_type, key);
-		/* gate and blackhole are together in a union */
-		key = jhash(&nh->gate, sizeof(nh->gate), key);
-		key = jhash(&nh->src, sizeof(nh->src), key);
-		key = jhash(&nh->rmap_src, sizeof(nh->rmap_src), key);
-		if (nh->nh_label) {
-			for (i = 0; i < nh->nh_label->num_labels; i++)
-				key = jhash_1word(nh->nh_label->label[i], key);
-		}
-		switch (nh->type) {
-		case NEXTHOP_TYPE_IPV4_IFINDEX:
-		case NEXTHOP_TYPE_IPV6_IFINDEX:
-		case NEXTHOP_TYPE_IFINDEX:
-			key = jhash_1word(nh->ifindex, key);
-			break;
-		case NEXTHOP_TYPE_BLACKHOLE:
-		case NEXTHOP_TYPE_IPV4:
-		case NEXTHOP_TYPE_IPV6:
-			break;
-		}
-	}
-	return key;
-}
-
 uint32_t zebra_nhg_hash_key(const void *arg)
 {
 	const struct nhg_hash_entry *nhe = arg;
@@ -298,7 +280,7 @@ uint32_t zebra_nhg_hash_key(const void *arg)
 
 	key = jhash_2words(nhe->vrf_id, nhe->afi, key);
 
-	key = jhash_1word(zebra_nhg_hash_key_nexthop_group(nhe->nhg), key);
+	key = jhash_1word(nexthop_group_hash(nhe->nhg), key);
 
 	return key;
 }
@@ -335,6 +317,31 @@ bool zebra_nhg_hash_id_equal(const void *arg1, const void *arg2)
 	return nhe1->id == nhe2->id;
 }
 
+/**
+ * zebra_nhg_cmp() - Compare the ID's of two nhe's
+ *
+ * @nhe1:	Nexthop group hash entry #1
+ * @nhe2:	Nexthop group hash entry #2
+ *
+ * Return:
+ * 	- Negative: #1 < #2
+ * 	- Positive: #1 > #2
+ * 	- Zero:	    #1 = #2
+ *
+ * This is used in the nhg RB trees.
+ */
+static int zebra_nhg_cmp(const struct nhg_hash_entry *nhe1,
+			 const struct nhg_hash_entry *nhe2)
+{
+	return nhe1->id - nhe2->id;
+}
+
+static int zebra_nhg_depend_cmp(const struct nhg_depend *dep1,
+				const struct nhg_depend *dep2)
+{
+	return zebra_nhg_cmp(dep1->nhe, dep2->nhe);
+}
+
 /**
  * zebra_nhg_find() - Find the zebra nhg in our table, or create it
  *
@@ -344,7 +351,7 @@ bool zebra_nhg_hash_id_equal(const void *arg1, const void *arg2)
  * @id:			ID we lookup with, 0 means its from us and we
  * 			need to give it an ID, otherwise its from the
  * 			kernel as we use the ID it gave us.
- * @nhg_depends:	Nexthop dependencies
+ * @nhg_depends:	Nexthop dependency tree head
  * @is_kernel_nh:	Was the nexthop created by the kernel
  *
  * Return:		Hash entry found or created
@@ -368,7 +375,7 @@ bool zebra_nhg_hash_id_equal(const void *arg1, const void *arg2)
  */
 struct nhg_hash_entry *zebra_nhg_find(struct nexthop_group *nhg,
 				      vrf_id_t vrf_id, afi_t afi, uint32_t id,
-				      struct list *nhg_depends,
+				      struct nhg_depends_head *nhg_depends,
 				      bool is_kernel_nh)
 {
 	/* lock for getiing and setting the id */
@@ -376,7 +383,7 @@ struct nhg_hash_entry *zebra_nhg_find(struct nexthop_group *nhg,
 	/* id counter to keep in sync with kernel */
 	static uint32_t id_counter = 0;
 
-	struct nhg_hash_entry lookup = {0};
+	struct nhg_hash_entry lookup = {};
 	struct nhg_hash_entry *nhe = NULL;
 	uint32_t old_id_counter = 0;
 
@@ -399,7 +406,7 @@ struct nhg_hash_entry *zebra_nhg_find(struct nexthop_group *nhg,
 	lookup.vrf_id = vrf_id;
 	lookup.afi = afi;
 	lookup.nhg = nhg;
-	lookup.nhg_depends = nhg_depends;
+	lookup.nhg_depends = *nhg_depends;
 	lookup.is_kernel_nh = is_kernel_nh;
 
 	if (id)
@@ -442,23 +449,34 @@ struct nhg_hash_entry *zebra_nhg_find_nexthop(struct nexthop *nh, afi_t afi)
 	return nhe;
 }
 
+void zebra_nhg_depends_free(struct nhg_depends_head *head)
+{
+	struct nhg_depend *rb_node_dep = NULL;
+	struct nhg_depend *tmp = NULL;
+
+	if (!zebra_nhg_depends_head_is_empty(head)) {
+		RB_FOREACH_SAFE (rb_node_dep, nhg_depends_head, head, tmp) {
+			RB_REMOVE(nhg_depends_head, head, rb_node_dep);
+			nhg_depend_free(rb_node_dep);
+		}
+	}
+}
+
 /**
  * zebra_nhg_free_group_depends() - Helper function for freeing nexthop_group
  * 				    struct and depends
  *
- * @nhg:		Nexthop group
- * @nhg_depends:	Nexthop group hash entry dependency list
+ * @nhg:		Nexthop_group
+ * @nhg_depends:	Nexthop group dependency tree head
  */
-void zebra_nhg_free_group_depends(struct nexthop_group *nhg,
-				  struct list *nhg_depends)
+void zebra_nhg_free_group_depends(struct nexthop_group **nhg,
+				  struct nhg_depends_head *head)
 {
-	if (nhg_depends)
-		list_delete(&nhg_depends);
-	if (nhg) {
-		if (nhg->nexthop)
-			nexthops_free(nhg->nexthop);
-		nexthop_group_delete(&nhg);
-	}
+	if (head)
+		zebra_nhg_depends_free(head);
+
+	if (nhg)
+		nexthop_group_free_delete(nhg);
 }
 
 /**
@@ -470,7 +488,7 @@ void zebra_nhg_free_group_depends(struct nexthop_group *nhg,
  */
 void zebra_nhg_free_members(struct nhg_hash_entry *nhe)
 {
-	zebra_nhg_free_group_depends(nhe->nhg, nhe->nhg_depends);
+	zebra_nhg_free_group_depends(&nhe->nhg, &nhe->nhg_depends);
 }
 
 /**
@@ -525,12 +543,11 @@ static void zebra_nhg_uninstall_release(struct nhg_hash_entry *nhe)
  */
 void zebra_nhg_decrement_ref(struct nhg_hash_entry *nhe)
 {
-	if (nhe->nhg_depends) {
-		struct listnode *ln = NULL;
-		struct nhg_depend *n_dp = NULL;
+	if (!zebra_nhg_depends_is_empty(nhe)) {
+		struct nhg_depend *rb_node_dep = NULL;
 
-		for (ALL_LIST_ELEMENTS_RO(nhe->nhg_depends, ln, n_dp)) {
-			zebra_nhg_decrement_ref(n_dp->nhe);
+		RB_FOREACH (rb_node_dep, nhg_depends_head, &nhe->nhg_depends) {
+			zebra_nhg_decrement_ref(rb_node_dep->nhe);
 		}
 	}
 
@@ -548,12 +565,11 @@ void zebra_nhg_decrement_ref(struct nhg_hash_entry *nhe)
  */
 void zebra_nhg_increment_ref(struct nhg_hash_entry *nhe)
 {
-	if (nhe->nhg_depends) {
-		struct listnode *ln = NULL;
-		struct nhg_depend *n_dp = NULL;
+	if (!zebra_nhg_depends_is_empty(nhe)) {
+		struct nhg_depend *rb_node_dep = NULL;
 
-		for (ALL_LIST_ELEMENTS_RO(nhe->nhg_depends, ln, n_dp)) {
-			zebra_nhg_increment_ref(n_dp->nhe);
+		RB_FOREACH (rb_node_dep, nhg_depends_head, &nhe->nhg_depends) {
+			zebra_nhg_increment_ref(rb_node_dep->nhe);
 		}
 	}
 
@@ -1177,7 +1193,7 @@ void zebra_nhg_dplane_result(struct zebra_dplane_ctx *ctx)
 	op = dplane_ctx_get_op(ctx);
 	status = dplane_ctx_get_status(ctx);
 
-	id = dplane_ctx_get_nhe(ctx)->id;
+	id = dplane_ctx_get_nhe_id(ctx);
 	nhe = zebra_nhg_lookup_id(id);
 
 	if (nhe) {
diff --git a/zebra/zebra_nhg.h b/zebra/zebra_nhg.h
index 90c7abb3e6..1cc1c47e32 100644
--- a/zebra/zebra_nhg.h
+++ b/zebra/zebra_nhg.h
@@ -28,6 +28,18 @@
 
 #include "zebra/zebra_dplane.h"
 
+/* This struct is used exclusively for dataplane
+ * interaction via a dataplane context.
+ *
+ * It is designed to mimic the netlink nexthop_grp
+ * struct in include/linux/nexthop.h
+ */
+struct depend_info {
+	uint32_t id;
+	uint8_t weight;
+};
+
+
 struct nhg_hash_entry {
 	uint32_t id;
 	afi_t afi;
@@ -49,12 +61,15 @@ struct nhg_hash_entry {
 
 	uint32_t flags;
 
-	/* Dependency list for other entries.
+	/* Dependency tree for other entries.
 	 * For instance a group with two
 	 * nexthops will have two dependencies
 	 * pointing to those nhg_hash_entries.
+	 *
+	 * Using a rb tree here to make lookups
+	 * faster with ID's.
 	 */
-	struct list *nhg_depends;
+	RB_HEAD(nhg_depends_head, nhg_depend) nhg_depends;
 /*
  * Is this nexthop group valid, ie all nexthops are fully resolved.
  * What is fully resolved?  It's a nexthop that is either self contained
@@ -75,22 +90,32 @@ struct nhg_hash_entry {
 #define NEXTHOP_GROUP_QUEUED 0x4
 };
 
-/* Struct for dependency nexthop */
+/* Abstraction for dependency tree */
 struct nhg_depend {
+	RB_ENTRY(nhg_depend) depend;
 	struct nhg_hash_entry *nhe;
 };
 
+RB_PROTOTYPE(nhg_depends_head, nhg_depend, depend, zebra_nhg_depend_cmp);
 
 void zebra_nhg_init(void);
 void zebra_nhg_terminate(void);
 
-extern struct nhg_depend *nhg_depend_add(struct list *nhg_depends,
-					 struct nhg_hash_entry *depend);
-extern struct nhg_depend *nhg_depend_new(void);
-extern void nhg_depend_free(struct nhg_depend *depends);
+extern uint8_t 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_depends_head_del(struct nhg_depends_head *head,
+				       struct nhg_hash_entry *depend);
+extern void zebra_nhg_depends_head_add(struct nhg_depends_head *head,
+				       struct nhg_hash_entry *depend);
+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_depends_head_init(struct nhg_depends_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);
 
-extern struct list *nhg_depend_new_list(void);
-extern struct list *nhg_depend_dup_list(struct list *from);
 
 extern struct nhg_hash_entry *zebra_nhg_lookup_id(uint32_t id);
 extern int zebra_nhg_insert_id(struct nhg_hash_entry *nhe);
@@ -103,13 +128,16 @@ extern bool zebra_nhg_hash_id_equal(const void *arg1, const void *arg2);
 
 extern struct nhg_hash_entry *
 zebra_nhg_find(struct nexthop_group *nhg, vrf_id_t vrf_id, afi_t afi,
-	       uint32_t id, struct list *nhg_depends, bool is_kernel_nh);
+	       uint32_t id, struct nhg_depends_head *nhg_depends,
+	       bool is_kernel_nh);
 
 extern struct nhg_hash_entry *zebra_nhg_find_nexthop(struct nexthop *nh,
 						     afi_t afi);
 
-void zebra_nhg_free_group_depends(struct nexthop_group *nhg,
-				  struct list *nhg_depends);
+
+void zebra_nhg_depends_free(struct nhg_depends_head *head);
+void zebra_nhg_free_group_depends(struct nexthop_group **nhg,
+				  struct nhg_depends_head *head);
 void zebra_nhg_free_members(struct nhg_hash_entry *nhe);
 void zebra_nhg_free(void *arg);
 void zebra_nhg_release(struct nhg_hash_entry *nhe);
@@ -123,5 +151,7 @@ void zebra_nhg_uninstall_kernel(struct nhg_hash_entry *nhe);
 
 void zebra_nhg_cleanup_tables(void);
 
+/* Forward ref of dplane update context type */
+struct zebra_dplane_ctx;
 void zebra_nhg_dplane_result(struct zebra_dplane_ctx *ctx);
 #endif
diff --git a/zebra/zebra_rib.c b/zebra/zebra_rib.c
index 3155570df5..957b7e14f4 100644
--- a/zebra/zebra_rib.c
+++ b/zebra/zebra_rib.c
@@ -2638,7 +2638,7 @@ int rib_add_multipath(afi_t afi, safi_t safi, struct prefix *p,
 	struct route_node *rn;
 	struct route_entry *same = NULL;
 	struct nhg_hash_entry *nhe = NULL;
-	struct list *nhg_depends = NULL;
+	struct nhg_depends_head nhg_depends = {0};
 	/* Default to route afi */
 	afi_t nhg_afi = afi;
 	int ret = 0;
@@ -2651,7 +2651,7 @@ int rib_add_multipath(afi_t afi, safi_t safi, struct prefix *p,
 	/* Lookup table.  */
 	table = zebra_vrf_table_with_table_id(afi, safi, re->vrf_id, re->table);
 	if (!table) {
-		zebra_nhg_free_group_depends(re->ng, nhg_depends);
+		nexthop_group_free_delete(&re->ng);
 		XFREE(MTYPE_RE, re);
 		return 0;
 	}
@@ -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;
 
-		nhg_depends = nhg_depend_new_list();
+		zebra_nhg_depends_head_init(&nhg_depends);
 
 		for (ALL_NEXTHOPS_PTR(re->ng, nh)) {
 			lookup = *nh;
@@ -2675,20 +2675,21 @@ 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);
-			nhg_depend_add(nhg_depends, depend);
+			zebra_nhg_depends_head_add(&nhg_depends, depend);
 		}
 
 		/* change the afi for group */
-		if (listcount(nhg_depends))
-			nhg_afi = AFI_UNSPEC;
+		nhg_afi = AFI_UNSPEC;
 	}
 
+	// TODO: Add proto type here
 	nhe = zebra_nhg_find(re->ng, re->vrf_id, nhg_afi, re->nhe_id,
-			     nhg_depends, false);
+			     &nhg_depends, false);
 
 	if (nhe) {
-		// TODO: Add interface pointer
-		zebra_nhg_free_group_depends(re->ng, nhg_depends);
+		/* It should point to the nhe nexthop group now */
+		if (re->ng)
+			nexthop_group_free_delete(&re->ng);
 		re->ng = nhe->nhg;
 		re->nhe_id = nhe->id;
 		zebra_nhg_increment_ref(nhe);
@@ -2697,6 +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_depends_free(&nhg_depends);
 	}
 
 
diff --git a/zebra/zebra_vty.c b/zebra/zebra_vty.c
index c2a9d091d2..106cd6b027 100644
--- a/zebra/zebra_vty.c
+++ b/zebra/zebra_vty.c
@@ -1127,13 +1127,13 @@ static void show_nexthop_group_cmd_helper(struct vty *vty,
 			vty_out(vty, "\tInterface Index: %d\n",
 				nhe->ifp->ifindex);
 
-		if (nhe->nhg_depends) {
-			struct listnode *dp_node = NULL;
-			struct nhg_depend *n_dp = NULL;
+		if (!zebra_nhg_depends_is_empty(nhe)) {
+			struct nhg_depend *rb_node_dep = NULL;
+
 			vty_out(vty, "\tDepends:");
-			for (ALL_LIST_ELEMENTS_RO(nhe->nhg_depends, dp_node,
-						  n_dp)) {
-				vty_out(vty, " (%u)", n_dp->nhe->id);
+			RB_FOREACH (rb_node_dep, nhg_depends_head,
+				    &nhe->nhg_depends) {
+				vty_out(vty, " (%u)", rb_node_dep->nhe->id);
 			}
 			vty_out(vty, "\n");
 		}
-- 
2.39.5