Protect homa_tx_pool_core with a local lock

homa_tx_pool.c

@@ -54,6 +54,7 @@ int homa_tx_pool_init(struct homa *homa)
 
 		nid = homa_tx_pool_mem_nid(i);
 		tx_core = &per_cpu(homa_tx_pool_core, i);
+		local_lock_init(&tx_core->bh_lock);
 		tx_core->nid = nid;
 	}
 	pr_notice("%s found max NUMA node %d\n", __func__, homa->max_numa);
@@ -172,6 +173,7 @@ int homa_tx_pool_alloc(struct homa *homa, int length, int *num_frags,
 	 * task preemption (so we can access per-core info without races).
 	 */
 	local_bh_disable();
+	local_lock_nested_bh(&homa_tx_pool_core.bh_lock);
 	tx_core = &per_cpu(homa_tx_pool_core, smp_processor_id());
 	tx_pool = homa->tx_pools[tx_core->nid];
 
@@ -216,6 +218,7 @@ int homa_tx_pool_alloc(struct homa *homa, int length, int *num_frags,
 		bytes_left -= skb_frag_size(frag);
 	}
 	nid = tx_core->nid;
+	local_unlock_nested_bh(&homa_tx_pool_core.bh_lock);
 	local_bh_enable();
 	*frags  = cur_frags;
 	*num_frags = frags_allocated;
@@ -226,6 +229,7 @@ int homa_tx_pool_alloc(struct homa *homa, int length, int *num_frags,
 	if (cur_frags != *frags)
 		kfree(cur_frags);
 	*num_frags = 0;
+	local_unlock_nested_bh(&homa_tx_pool_core.bh_lock);
 	local_bh_enable();
 	return err;
 }

homa_tx_pool.h

@@ -56,6 +56,13 @@ struct homa_tx_pool {
  * tx memory allocation. All values are assumed to be zero initially.
  */
 struct homa_tx_pool_core {
+	/**
+	 * @bh_lock: protects @page, @page_size, and @allocated. Disabling
+	 * BHs is not enough on PREEMPT_RT, where another task can run on
+	 * this core while BHs are disabled.
+	 */
+	local_lock_t bh_lock;
+
 	/**
 	 * @page: a page of data that is "owned" by this core and used for
 	 * allocation on this core (or NULL). This pointer is included in

test/mock.c

@@ -210,6 +210,11 @@ int mock_rpc_holds;
  */
 static int mock_local_bh_disables;
 
+/* The number of times local_lock_nested_bh() has been invoked, minus the
+ * number of times local_unlock_nested_bh has been invoked.
+ */
+static int mock_local_bh_locks;
+
 /* The number of times preempt_disable() has been invoked, minus the
  * number of times preempt_enable has been invoked.
  */
@@ -2245,6 +2250,20 @@ void mock_local_bh_enable(void)
 	mock_local_bh_disables--;
 }
 
+void mock_local_lock_nested_bh(void)
+{
+	if (mock_local_bh_disables == 0)
+		FAIL(" local_lock_nested_bh invoked without local_bh_disable");
+	mock_local_bh_locks++;
+}
+
+void mock_local_unlock_nested_bh(void)
+{
+	if (mock_local_bh_locks == 0)
+		FAIL(" local_unlock_nested_bh invoked without local_lock_nested_bh");
+	mock_local_bh_locks--;
+}
+
 /**
  * mock_net_for_hnet() - Return the struct net associated with a struct
  * homa_net, or NULL if the struct net can't be identified.
@@ -3034,6 +3053,11 @@ void mock_teardown(void)
 				mock_local_bh_disables);
 	mock_local_bh_disables = 0;
 
+	if (mock_local_bh_locks != 0)
+		FAIL(" %d local_lock_nested_bhs still active after test",
+				mock_local_bh_locks);
+	mock_local_bh_locks = 0;
+
 	if (mock_preempt_disables != 0)
 		FAIL(" %d preempt_disables still active after test",
 				mock_preempt_disables);

test/mock.h

@@ -87,6 +87,12 @@
 #undef local_irq_save
 #define local_irq_save(flags) (flags) = 0
 
+#undef local_lock_nested_bh
+#define local_lock_nested_bh(lock) mock_local_lock_nested_bh()
+
+#undef local_unlock_nested_bh
+#define local_unlock_nested_bh(lock) mock_local_unlock_nested_bh()
+
 #undef MAX_SKB_FRAGS
 #define MAX_SKB_FRAGS mock_max_skb_frags
 
@@ -295,6 +301,8 @@ bool        mock_is_locked(void *lock);
 void       *mock_kmalloc(size_t size, gfp_t flags);
 void        mock_local_bh_disable(void);
 void        mock_local_bh_enable(void);
+void        mock_local_lock_nested_bh(void);
+void        mock_local_unlock_nested_bh(void);
 struct net *mock_net_for_hnet(struct homa_net *hnet);
 void       *mock_net_generic(const struct net *net, unsigned int id);
 int         mock_page_refs(struct page *page);