bsc_subscriber: Optimize lookup of bsub by TMSI

It was found that on a busy osmo-bsc process (>1000 concurrent calls
spead over different BTSs), a good amount of time is spent iterating the
subscribers list trying to find a subscriber based on a TMSI (1.60% of
total CPU time used by osmo-bsc).

This patch introduces a new rbtree under struct bsc_subscr_store which
allows storing all the busbs ordered by TMSI.
This way, lookup time changes O(N) -> O(log(N)), at the expense of
increased insert/deletion time O(1) -> O(log(N)).

Related: SYS#6200
Change-Id: If27429e715ef4b327177e427249e68321a6e83cc

4 files changed, 84 insertions, 6 deletions

diff --git a/include/osmocom/bsc/bsc_subscriber.h b/include/osmocom/bsc/bsc_subscriber.h
index 7ec457bf4..1fa00162d 100644
--- a/include/osmocom/bsc/bsc_subscriber.h
+++ b/include/osmocom/bsc/bsc_subscriber.h
@@ -5,6 +5,7 @@
 #include <stdint.h>
 
 #include <osmocom/core/linuxlist.h>
+#include <osmocom/core/linuxrbtree.h>
 #include <osmocom/core/use_count.h>
 #include <osmocom/gsm/protocol/gsm_23_003.h>
 #include <osmocom/gsm/gsm48.h>
@@ -14,6 +15,8 @@ struct gsm_bts;
 
 struct bsc_subscr_store {
 	struct llist_head bsub_list; /* list containing "struct bsc_subscr" */
+	/* rbtree root of 'struct bsc_subscr', ordered by tmsi */
+	struct rb_root bsub_tree_tmsi;
 };
 
 struct bsc_subscr_store *bsc_subscr_store_alloc(void *ctx);
@@ -21,6 +24,8 @@ struct bsc_subscr_store *bsc_subscr_store_alloc(void *ctx);
 struct bsc_subscr {
 	struct bsc_subscr_store *store; /* backpointer to "struct bsc_subscr_store" */
 	struct llist_head entry; /* entry in (struct bsc_subscr_store)->bsub_list */
+	/* entry in (struct bsc_subscr_store)->bsub_tree_tmsi. Inserted if tmsi != GSM_RESERVED_TMSI: */
+	struct rb_node node_tmsi;
 	struct osmo_use_count use_count;
 
 	char imsi[GSM23003_IMSI_MAX_DIGITS+1];
@@ -49,6 +54,7 @@ struct bsc_subscr *bsc_subscr_find_by_imsi(struct bsc_subscr_store *bsubst,
 					   const char *imsi,
 					   const char *use_token);
 
+int bsc_subscr_set_tmsi(struct bsc_subscr *bsub, uint32_t tmsi);
 void bsc_subscr_set_imsi(struct bsc_subscr *bsub, const char *imsi);
 void bsc_subscr_set_imei(struct bsc_subscr *bsub, const char *imei);
 
diff --git a/src/osmo-bsc/bsc_subscriber.c b/src/osmo-bsc/bsc_subscriber.c
index 65c5c8764..1f69d442d 100644
--- a/src/osmo-bsc/bsc_subscriber.c
+++ b/src/osmo-bsc/bsc_subscriber.c
@@ -138,20 +138,81 @@ static struct bsc_subscr *bsc_subscr_find_by_tmsi(struct bsc_subscr_store *bsubs
 						  uint32_t tmsi,
 						  const char *use_token)
 {
+	const struct rb_node *node = bsubst->bsub_tree_tmsi.rb_node;
 	struct bsc_subscr *bsub;
 
 	if (tmsi == GSM_RESERVED_TMSI)
 		return NULL;
 
-	llist_for_each_entry(bsub, &bsubst->bsub_list, entry) {
-		if (bsub->tmsi == tmsi) {
+	while (node) {
+		bsub = container_of(node, struct bsc_subscr, node_tmsi);
+		if (tmsi < bsub->tmsi)
+			node = node->rb_left;
+		else if (tmsi > bsub->tmsi)
+			node = node->rb_right;
+		else {
 			bsc_subscr_get(bsub, use_token);
 			return bsub;
 		}
 	}
+
 	return NULL;
 }
 
+static int bsc_subscr_store_insert_bsub_tmsi(struct bsc_subscr *bsub)
+{
+	struct bsc_subscr_store *bsubst = bsub->store;
+	struct rb_node **n = &(bsubst->bsub_tree_tmsi.rb_node);
+	struct rb_node *parent = NULL;
+
+	OSMO_ASSERT(bsub->tmsi != GSM_RESERVED_TMSI);
+
+	while (*n) {
+		struct bsc_subscr *it;
+
+		it = container_of(*n, struct bsc_subscr, node_tmsi);
+
+		parent = *n;
+		if (bsub->tmsi < it->tmsi) {
+			n = &((*n)->rb_left);
+		} else if (bsub->tmsi > it->tmsi) {
+			n = &((*n)->rb_right);
+		} else {
+			LOGP(DMSC, LOGL_ERROR, "Trying to reserve already reserved tmsi %u\n", bsub->tmsi);
+			return -EEXIST;
+		}
+	}
+
+	rb_link_node(&bsub->node_tmsi, parent, n);
+	rb_insert_color(&bsub->node_tmsi, &bsubst->bsub_tree_tmsi);
+	return 0;
+}
+
+int bsc_subscr_set_tmsi(struct bsc_subscr *bsub, uint32_t tmsi)
+{
+	int rc = 0;
+
+	if (!bsub)
+		return -EINVAL;
+
+	if (bsub->tmsi == tmsi)
+		return 0;
+
+	/* bsub was already inserted, remove and re-insert with new tmsi */
+	if (bsub->tmsi != GSM_RESERVED_TMSI)
+		rb_erase(&bsub->node_tmsi, &bsub->store->bsub_tree_tmsi);
+
+	bsub->tmsi = tmsi;
+
+	/* If new tmsi is set, insert bsub into rbtree: */
+	if (bsub->tmsi != GSM_RESERVED_TMSI) {
+		if ((rc = bsc_subscr_store_insert_bsub_tmsi(bsub)) < 0)
+			bsub->tmsi = GSM_RESERVED_TMSI;
+	}
+
+	return rc;
+}
+
 void bsc_subscr_set_imsi(struct bsc_subscr *bsub, const char *imsi)
 {
 	if (!bsub)
@@ -209,7 +270,10 @@ struct bsc_subscr *bsc_subscr_find_or_create_by_tmsi(struct bsc_subscr_store *bs
 	bsub = bsc_subscr_alloc(bsubst);
 	if (!bsub)
 		return NULL;
-	bsub->tmsi = tmsi;
+	if (bsc_subscr_set_tmsi(bsub, tmsi) < 0) {
+		bsc_subscr_free(bsub);
+		return NULL;
+	}
 	bsc_subscr_get(bsub, use_token);
 	return bsub;
 }
@@ -268,6 +332,10 @@ const char *bsc_subscr_id(struct bsc_subscr *bsub)
 static void bsc_subscr_free(struct bsc_subscr *bsub)
 {
 	OSMO_ASSERT(llist_empty(&bsub->active_paging_requests));
+
+	if (bsub->tmsi != GSM_RESERVED_TMSI)
+		rb_erase(&bsub->node_tmsi, &bsub->store->bsub_tree_tmsi);
+
 	llist_del(&bsub->entry);
 	talloc_free(bsub);
 }
diff --git a/src/osmo-bsc/osmo_bsc_bssap.c b/src/osmo-bsc/osmo_bsc_bssap.c
index a7c6ac355..d504ce691 100644
--- a/src/osmo-bsc/osmo_bsc_bssap.c
+++ b/src/osmo-bsc/osmo_bsc_bssap.c
@@ -352,8 +352,12 @@ int bsc_paging_start(struct bsc_paging_params *params)
 			return -EINVAL;
 		}
 	}
-	if (params->tmsi != GSM_RESERVED_TMSI)
-		params->bsub->tmsi = params->tmsi;
+	if (params->tmsi != GSM_RESERVED_TMSI) {
+		if (bsc_subscr_set_tmsi(params->bsub, params->tmsi) < 0) {
+			LOG_PAGING(params, DMSC, LOGL_ERROR, "Paging request failed: Could not set TMSI on subscriber\n");
+			return -EINVAL;
+		}
+	}
 	log_set_context(LOG_CTX_BSC_SUBSCR, params->bsub);
 
 	switch (params->cil.id_discr) {
diff --git a/tests/subscr/bsc_subscr_test.c b/tests/subscr/bsc_subscr_test.c
index 27dc93e1d..2a0e1eb0c 100644
--- a/tests/subscr/bsc_subscr_test.c
+++ b/tests/subscr/bsc_subscr_test.c
@@ -75,7 +75,7 @@ static void test_bsc_subscr(void)
 
 	/* Allocate entry 2 */
 	s2 = bsc_subscr_find_or_create_by_imsi(bsc_subscribers, imsi2, BSUB_USE);
-	s2->tmsi = 0x73517351;
+	bsc_subscr_set_tmsi(s2, 0x73517351);
 	VERBOSE_ASSERT(llist_count(&bsc_subscribers->bsub_list), == 2, "%d");
 
 	/* Allocate entry 3 */