diff --git a/tools/testing/selftests/bpf/prog_tests/test_xsk.c b/tools/testing/selftests/bpf/prog_tests/test_xsk.c
index 4549358cc8c21..5d4eb47cbfee5 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_xsk.c
+++ b/tools/testing/selftests/bpf/prog_tests/test_xsk.c
@@ -101,6 +101,7 @@ int xsk_configure_umem(struct ifobject *ifobj, struct xsk_umem_info *umem, void
 		return ret;
 
 	umem->buffer = buffer;
+	refcount_set(&umem->users, 1);
 	if (ifobj->shared_umem && ifobj->rx_on) {
 		umem->base_addr = umem_size(umem);
 		umem->next_buffer = umem_size(umem);
@@ -154,6 +155,7 @@ int xsk_configure_socket(struct xsk_socket_info *xsk, struct xsk_umem_info *umem
 	struct xsk_socket_config cfg = {};
 	struct xsk_ring_cons *rxr;
 	struct xsk_ring_prod *txr;
+	int ret;
 
 	xsk->umem = umem;
 	cfg.rx_size = xsk->rxqsize;
@@ -170,7 +172,26 @@ int xsk_configure_socket(struct xsk_socket_info *xsk, struct xsk_umem_info *umem
 
 	txr = ifobject->tx_on ? &xsk->tx : NULL;
 	rxr = ifobject->rx_on ? &xsk->rx : NULL;
-	return xsk_socket__create(&xsk->xsk, ifobject->ifindex, 0, umem->umem, rxr, txr, &cfg);
+	ret = xsk_socket__create(&xsk->xsk, ifobject->ifindex, 0, umem->umem, rxr, txr, &cfg);
+	if (ret) {
+		/*
+		 * For shared sockets refcount_inc hasn't run yet, so clear umem to mark this slot
+		 * as having no reference. For the owner (non-shared) the reference was taken by
+		 * xsk_configure_umem; leave umem set so the caller's rollback path can release it
+		 * via umem_ref.
+		 */
+		if (shared)
+			xsk->umem = NULL;
+		xsk->xsk = NULL;
+		return ret;
+	}
+
+	if (shared) {
+		refcount_inc(&umem->users);
+		xsk->umem_ref = true;
+	}
+
+	return ret;
 }
 
 static int set_ring_size(struct ifobject *ifobj)
@@ -203,6 +224,12 @@ int hw_ring_size_reset(struct ifobject *ifobj)
 static void __test_spec_init(struct test_spec *test, struct ifobject *ifobj_tx,
 			     struct ifobject *ifobj_rx)
 {
+	/*
+	 * Keep the same default as xskxceiver startup: when TX and RX share the same netdev,
+	 * shared UMEM is the baseline mode for this test harness. Individual tests can still
+	 * override this as needed.
+	 */
+	bool shared_default = ifobj_tx->ifindex == ifobj_rx->ifindex;
 	u32 i, j;
 
 	for (i = 0; i < MAX_INTERFACES; i++) {
@@ -214,6 +241,7 @@ static void __test_spec_init(struct test_spec *test, struct ifobject *ifobj_tx,
 		ifobj->use_fill_ring = true;
 		ifobj->release_rx = true;
 		ifobj->validation_func = NULL;
+		ifobj->shared_umem = shared_default;
 		ifobj->use_metadata = false;
 
 		if (i == 0) {
@@ -572,28 +600,95 @@ static int pkt_stream_receive_half(struct test_spec *test)
 	return 0;
 }
 
-static int pkt_stream_even_odd_sequence(struct test_spec *test)
+/*
+ * Regenerate every per-socket stream, publishing into xsk_arr[]
+ * only once all allocations have succeeded. @dims supplies the
+ * per-socket packet count and length, derived from the stream
+ * currently installed in the slot.
+ *
+ * The replaced streams are the defaults owned by test_spec, which
+ * every slot aliases, so they must not be freed here.
+ */
+static int pkt_stream_replace_seq(struct test_spec *test, pkt_stream_dims_fn dims,
+				  const void *ctx)
 {
-	struct pkt_stream *pkt_stream;
-	u32 i;
+	struct pkt_stream *tx_streams[MAX_SOCKETS] = {};
+	struct pkt_stream *rx_streams[MAX_SOCKETS] = {};
+	u32 i, nb_pkts, pkt_len;
 
 	for (i = 0; i < test->nb_sockets; i++) {
-		pkt_stream = test->ifobj_tx->xsk_arr[i].pkt_stream;
-		pkt_stream = __pkt_stream_generate(pkt_stream->nb_pkts / 2,
-						   pkt_stream->pkts[0].len, i, 2);
-		if (!pkt_stream)
-			return -ENOMEM;
-		test->ifobj_tx->xsk_arr[i].pkt_stream = pkt_stream;
+		dims(test->ifobj_tx->xsk_arr[i].pkt_stream, i, ctx, &nb_pkts, &pkt_len);
+		tx_streams[i] = __pkt_stream_generate(nb_pkts, pkt_len, i, 2);
+		if (!tx_streams[i])
+			goto err;
+
+		dims(test->ifobj_rx->xsk_arr[i].pkt_stream, i, ctx, &nb_pkts, &pkt_len);
+		rx_streams[i] = __pkt_stream_generate(nb_pkts, pkt_len, i, 2);
+		if (!rx_streams[i])
+			goto err;
+	}
 
-		pkt_stream = test->ifobj_rx->xsk_arr[i].pkt_stream;
-		pkt_stream = __pkt_stream_generate(pkt_stream->nb_pkts / 2,
-						   pkt_stream->pkts[0].len, i, 2);
-		if (!pkt_stream)
-			return -ENOMEM;
-		test->ifobj_rx->xsk_arr[i].pkt_stream = pkt_stream;
+	for (i = 0; i < test->nb_sockets; i++) {
+		test->ifobj_tx->xsk_arr[i].pkt_stream = tx_streams[i];
+		test->ifobj_rx->xsk_arr[i].pkt_stream = rx_streams[i];
 	}
 
 	return 0;
+
+err:
+	for (i = 0; i < test->nb_sockets; i++) {
+		if (tx_streams[i])
+			pkt_stream_delete(tx_streams[i]);
+		if (rx_streams[i])
+			pkt_stream_delete(rx_streams[i]);
+	}
+
+	return -ENOMEM;
+}
+
+static void even_odd_dims(struct pkt_stream *orig, u32 sock_id, const void *ctx,
+			  u32 *nb_pkts, u32 *pkt_len)
+{
+	*nb_pkts = orig->nb_pkts / 2;
+	*pkt_len = orig->pkts[0].len;
+}
+
+static void len_dims(struct pkt_stream *orig, u32 sock_id, const void *ctx,
+		     u32 *nb_pkts, u32 *pkt_len)
+{
+	const struct shared_umem_len_ctx *cfg = ctx;
+
+	*nb_pkts = orig->nb_pkts / 2;
+	*pkt_len = sock_id ? cfg->long_len : cfg->short_len;
+}
+
+static void uneven_dist_dims(struct pkt_stream *orig, u32 sock_id, const void *ctx,
+			     u32 *nb_pkts, u32 *pkt_len)
+{
+	const struct shared_umem_uneven_dist_ctx *cfg = ctx;
+	u32 pkts_sock0 = cfg->total_pkts / 4;
+
+	*nb_pkts = sock_id ? cfg->total_pkts - pkts_sock0 : pkts_sock0;
+	*pkt_len = cfg->pkt_len;
+}
+
+static int pkt_stream_even_odd_sequence(struct test_spec *test)
+{
+	return pkt_stream_replace_seq(test, even_odd_dims, NULL);
+}
+
+static int pkt_stream_len_seq(struct test_spec *test, const struct shared_umem_len_ctx *cfg)
+{
+	return pkt_stream_replace_seq(test, len_dims, cfg);
+}
+
+static int pkt_stream_uneven_dist_seq(struct test_spec *test,
+				      const struct shared_umem_uneven_dist_ctx *cfg)
+{
+	if (test->nb_sockets < 2 || cfg->total_pkts < 4)
+		return -EINVAL;
+
+	return pkt_stream_replace_seq(test, uneven_dist_dims, cfg);
 }
 
 static void release_even_odd_sequence(struct test_spec *test)
@@ -1508,8 +1603,15 @@ static int thread_common_ops_tx(struct test_spec *test, struct ifobject *ifobjec
 	}
 
 	umem_rx = test->ifobj_rx->xsk_arr[0].umem;
+	/* Non-owning view used only for TX buffer arithmetic; the sockets below bind to
+	 * the RX-owned UMEM, so the handle, cq and refcount stay in one place.
+	 */
 	umem_tx = ifobject->xsk_arr[0].umem_real;
-	memcpy(umem_tx, umem_rx, sizeof(*umem_tx));
+	umem_tx->num_frames = umem_rx->num_frames;
+	umem_tx->frame_headroom = umem_rx->frame_headroom;
+	umem_tx->buffer = umem_rx->buffer;
+	umem_tx->frame_size = umem_rx->frame_size;
+	umem_tx->unaligned_mode = umem_rx->unaligned_mode;
 	umem_tx->base_addr = 0;
 	umem_tx->next_buffer = 0;
 
@@ -1599,6 +1701,8 @@ static int thread_common_ops(struct test_spec *test, struct ifobject *ifobject)
 	ret = xsk_configure_umem(ifobject, umem, bufs, umem_sz);
 	if (ret)
 		return ret;
+	/* Mark before xsk_configure so rollback can release the UMEM if it fails. */
+	ifobject->xsk->umem_ref = true;
 
 	ret = xsk_configure(test, ifobject, umem, false);
 	if (ret)
@@ -1669,12 +1773,15 @@ void *worker_testapp_validate_rx(void *arg)
 				       strerror(-err));
 	}
 
+	/* Publish setup failure before releasing the main thread from the barrier. */
+	if (err)
+		test->fail = true;
+
 	if (test->use_barrier)
 		pthread_barrier_wait(&barr);
 
 	/* We leave only now in case of error to avoid getting stuck in the barrier */
 	if (err) {
-		test->fail = true;
 		pthread_exit(NULL);
 	}
 
@@ -1701,12 +1808,65 @@ void *worker_testapp_validate_rx(void *arg)
 	pthread_exit(NULL);
 }
 
-static void testapp_clean_xsk_umem(struct ifobject *ifobj)
+int xsk_delete_socket(struct xsk_socket_info *xsk)
 {
-	struct xsk_umem_info *umem = ifobj->xsk->umem;
+	struct xsk_umem_info *umem = xsk->umem;
+
+	if (!umem)
+		return 0;
+
+	if (xsk->xsk)
+		xsk_socket__delete(xsk->xsk);
+	xsk->xsk = NULL;
 
-	xsk_umem__delete(umem->umem);
-	munmap(umem->buffer, umem->mmap_size);
+	/* Skip slots that never acquired a UMEM reference (pre-initialized but unconfigured). */
+	if (!xsk->umem_ref) {
+		xsk->umem = NULL;
+		return 0;
+	}
+
+	if (refcount_dec_and_test(&umem->users)) {
+		if (umem->umem) {
+			int err = xsk_umem__delete(umem->umem);
+
+			if (err) {
+				ksft_print_msg("xsk_umem__delete failed: %d (umem still busy?)\n",
+					       err);
+				/* Keep ownership explicit so a later cleanup pass can retry
+				 * delete.
+				 */
+				refcount_set(&umem->users, 1);
+				xsk->umem_ref = true;
+				xsk->umem = umem;
+				return err;
+			}
+			umem->umem = NULL;
+		}
+		if (umem->buffer && umem->mmap_size) {
+			munmap(umem->buffer, umem->mmap_size);
+			umem->buffer = NULL;
+			umem->mmap_size = 0;
+		}
+	}
+
+	xsk->umem_ref = false;
+	xsk->umem = NULL;
+	return 0;
+}
+
+static void xsk_delete_all_ifobj_sockets(struct test_spec *test, struct ifobject *ifobj)
+{
+	u32 i;
+
+	if (!ifobj)
+		return;
+
+	/* A UMEM that could not be deleted stays mapped until the process exits,
+	 * so record it rather than letting teardown drop the error.
+	 */
+	for (i = test->nb_sockets; i > 0; i--)
+		if (xsk_delete_socket(&ifobj->xsk_arr[i - 1]))
+			test->fail = true;
 }
 
 static bool xdp_prog_changed_rx(struct test_spec *test)
@@ -1768,27 +1928,6 @@ static int xsk_attach_xdp_progs(struct test_spec *test, struct ifobject *ifobj_r
 	return err;
 }
 
-static void clean_sockets(struct test_spec *test, struct ifobject *ifobj)
-{
-	u32 i;
-
-	if (!ifobj || !test)
-		return;
-
-	for (i = 0; i < test->nb_sockets; i++)
-		xsk_socket__delete(ifobj->xsk_arr[i].xsk);
-}
-
-static void clean_umem(struct test_spec *test, struct ifobject *ifobj1, struct ifobject *ifobj2)
-{
-	if (!ifobj1)
-		return;
-
-	testapp_clean_xsk_umem(ifobj1);
-	if (ifobj2 && !ifobj2->shared_umem)
-		testapp_clean_xsk_umem(ifobj2);
-}
-
 static int __testapp_validate_traffic(struct test_spec *test, struct ifobject *ifobj1,
 				      struct ifobject *ifobj2)
 {
@@ -1840,13 +1979,12 @@ static int __testapp_validate_traffic(struct test_spec *test, struct ifobject *i
 		if (pthread_barrier_destroy(&barr)) {
 			test->use_barrier = false;
 			pthread_join(t0, NULL);
-			clean_sockets(test, ifobj1);
-			clean_umem(test, ifobj1, NULL);
+			xsk_delete_all_ifobj_sockets(test, ifobj1);
 			return TEST_FAILURE;
 		}
 	}
 
-	if (ifobj2) {
+	if (ifobj2 && !test->fail) {
 		/*Spawn TX thread */
 		pthread_create(&t1, NULL, ifobj2->func_ptr, test);
 		pthread_join(t1, NULL);
@@ -1855,9 +1993,8 @@ static int __testapp_validate_traffic(struct test_spec *test, struct ifobject *i
 	pthread_join(t0, NULL);
 
 	if (test->total_steps == test->current_step || test->fail) {
-		clean_sockets(test, ifobj1);
-		clean_sockets(test, ifobj2);
-		clean_umem(test, ifobj1, ifobj2);
+		xsk_delete_all_ifobj_sockets(test, ifobj2);
+		xsk_delete_all_ifobj_sockets(test, ifobj1);
 	}
 
 	if (test->fail)
@@ -1966,9 +2103,8 @@ int testapp_xdp_prog_cleanup(struct test_spec *test)
 		return TEST_FAILURE;
 
 	if (swap_xsk_resources(test)) {
-		clean_sockets(test, test->ifobj_rx);
-		clean_sockets(test, test->ifobj_tx);
-		clean_umem(test, test->ifobj_rx, test->ifobj_tx);
+		xsk_delete_all_ifobj_sockets(test, test->ifobj_tx);
+		xsk_delete_all_ifobj_sockets(test, test->ifobj_rx);
 		return TEST_FAILURE;
 	}
 
@@ -2226,6 +2362,101 @@ int testapp_xdp_shared_umem(struct test_spec *test)
 	return ret;
 }
 
+static int shared_umem_seq_even_odd(struct test_spec *test, const void *ctx)
+{
+	return pkt_stream_even_odd_sequence(test) ? TEST_FAILURE : TEST_PASS;
+}
+
+static int shared_umem_seq_len(struct test_spec *test, const void *ctx)
+{
+	return pkt_stream_len_seq(test, ctx) ? TEST_FAILURE : TEST_PASS;
+}
+
+static int shared_umem_seq_uneven_dist(struct test_spec *test, const void *ctx)
+{
+	return pkt_stream_uneven_dist_seq(test, ctx) ? TEST_FAILURE : TEST_PASS;
+}
+
+static int run_shared_umem_test(struct test_spec *test, struct bpf_program *xdp_prog_rx,
+				struct bpf_program *xdp_prog_tx, struct bpf_map *xskmap_rx,
+				struct bpf_map *xskmap_tx, u32 nb_sockets,
+				shared_umem_seq_fn seq_fn, const void *ctx)
+{
+	int ret;
+
+	if (nb_sockets > MAX_SOCKETS) {
+		ksft_print_msg("ERROR: [%s] invalid socket count %u\n", __func__, nb_sockets);
+		return TEST_FAILURE;
+	}
+
+	test->total_steps = 1;
+	test->nb_sockets = nb_sockets;
+
+	test_spec_set_xdp_prog(test, xdp_prog_rx, xdp_prog_tx, xskmap_rx, xskmap_tx);
+
+	ret = seq_fn(test, ctx);
+	if (ret)
+		return ret;
+
+	ret = testapp_validate_traffic(test);
+
+	release_even_odd_sequence(test);
+
+	return ret;
+}
+
+int testapp_shared_umem_4_sockets(struct test_spec *test)
+{
+	struct xsk_xdp_progs *skel_rx = test->ifobj_rx->xdp_progs;
+	struct xsk_xdp_progs *skel_tx = test->ifobj_tx->xdp_progs;
+
+	return run_shared_umem_test(test, skel_rx->progs.xsk_xdp_shared_umem,
+				    skel_tx->progs.xsk_xdp_shared_umem, skel_rx->maps.xsk,
+				    skel_tx->maps.xsk, 4, shared_umem_seq_even_odd, NULL);
+}
+
+int testapp_shared_umem_length_based(struct test_spec *test)
+{
+	struct xsk_xdp_progs *skel_rx = test->ifobj_rx->xdp_progs;
+	struct xsk_xdp_progs *skel_tx = test->ifobj_tx->xdp_progs;
+	const struct shared_umem_len_ctx len_ctx = {
+		.short_len = MIN_PKT_SIZE,
+		.long_len = MIN_PKT_SIZE * 2,
+	};
+
+	return run_shared_umem_test(test, skel_rx->progs.xsk_xdp_shared_umem_length_based,
+				    skel_tx->progs.xsk_xdp_shared_umem_length_based,
+				    skel_rx->maps.xsk, skel_tx->maps.xsk, 2, shared_umem_seq_len,
+				    &len_ctx);
+}
+
+int testapp_shared_umem_uneven_dist(struct test_spec *test)
+{
+	struct xsk_xdp_progs *skel_rx = test->ifobj_rx->xdp_progs;
+	struct xsk_xdp_progs *skel_tx = test->ifobj_tx->xdp_progs;
+	const struct shared_umem_uneven_dist_ctx uneven_dist_ctx = {
+		.total_pkts = DEFAULT_PKT_CNT * 4,
+		.pkt_len = MIN_PKT_SIZE,
+	};
+
+	return run_shared_umem_test(test, skel_rx->progs.xsk_xdp_shared_umem,
+				    skel_tx->progs.xsk_xdp_shared_umem, skel_rx->maps.xsk,
+				    skel_tx->maps.xsk, 2, shared_umem_seq_uneven_dist,
+				    &uneven_dist_ctx);
+}
+
+int testapp_shared_umem_unaligned(struct test_spec *test)
+{
+	struct xsk_xdp_progs *skel_rx = test->ifobj_rx->xdp_progs;
+	struct xsk_xdp_progs *skel_tx = test->ifobj_tx->xdp_progs;
+
+	test_spec_set_unaligned(test);
+
+	return run_shared_umem_test(test, skel_rx->progs.xsk_xdp_shared_umem,
+				    skel_tx->progs.xsk_xdp_shared_umem, skel_rx->maps.xsk,
+				    skel_tx->maps.xsk, 2, shared_umem_seq_even_odd, NULL);
+}
+
 int testapp_poll_txq_tmout(struct test_spec *test)
 {
 	bool shared_umem = test->ifobj_tx->shared_umem;
@@ -2506,9 +2737,8 @@ int testapp_hw_sw_max_ring_size(struct test_spec *test)
 	test->ifobj_tx->xsk->batch_size = test->ifobj_tx->ring.tx_max_pending - 8;
 	test->ifobj_rx->xsk->batch_size = test->ifobj_tx->ring.tx_max_pending - 8;
 	if (pkt_stream_replace(test, max_descs, MIN_PKT_SIZE)) {
-		clean_sockets(test, test->ifobj_tx);
-		clean_sockets(test, test->ifobj_rx);
-		clean_umem(test, test->ifobj_rx, test->ifobj_tx);
+		xsk_delete_all_ifobj_sockets(test, test->ifobj_tx);
+		xsk_delete_all_ifobj_sockets(test, test->ifobj_rx);
 		return TEST_FAILURE;
 	}
 
diff --git a/tools/testing/selftests/bpf/prog_tests/test_xsk.h b/tools/testing/selftests/bpf/prog_tests/test_xsk.h
index 03753ddc5dcda..94c77574f4a8b 100644
--- a/tools/testing/selftests/bpf/prog_tests/test_xsk.h
+++ b/tools/testing/selftests/bpf/prog_tests/test_xsk.h
@@ -4,6 +4,7 @@
 
 #include <linux/ethtool.h>
 #include <linux/if_xdp.h>
+#include <linux/refcount.h>
 
 #include "../kselftest.h"
 #include "xsk.h"
@@ -76,9 +77,13 @@ enum test_mode {
 
 struct ifobject;
 struct test_spec;
+struct pkt_stream;
 typedef int (*validation_func_t)(struct ifobject *ifobj);
 typedef void *(*thread_func_t)(void *arg);
 typedef int (*test_func_t)(struct test_spec *test);
+typedef int (*shared_umem_seq_fn)(struct test_spec *test, const void *ctx);
+typedef void (*pkt_stream_dims_fn)(struct pkt_stream *orig, u32 sock_id, const void *ctx,
+				   u32 *nb_pkts, u32 *pkt_len);
 
 struct xsk_socket_info {
 	struct xsk_ring_cons rx;
@@ -93,6 +98,7 @@ struct xsk_socket_info {
 	u8 dst_mac[ETH_ALEN];
 	u8 src_mac[ETH_ALEN];
 	bool check_consumer;
+	bool umem_ref; /* true if this slot holds a counted UMEM reference */
 };
 
 int kick_rx(struct xsk_socket_info *xsk);
@@ -104,6 +110,7 @@ struct xsk_umem_info {
 	struct xsk_umem *umem;
 	u64 next_buffer;
 	u64 mmap_size;
+	refcount_t users;
 	u32 num_frames;
 	u32 frame_headroom;
 	void *buffer;
@@ -159,7 +166,7 @@ int init_iface(struct ifobject *ifobj, thread_func_t func_ptr);
 int xsk_configure_umem(struct ifobject *ifobj, struct xsk_umem_info *umem, void *buffer, u64 size);
 int xsk_configure_socket(struct xsk_socket_info *xsk, struct xsk_umem_info *umem,
 			 struct ifobject *ifobject, bool shared);
-
+int xsk_delete_socket(struct xsk_socket_info *xsk);
 
 struct pkt {
 	int offset;
@@ -179,6 +186,16 @@ struct pkt_stream {
 	bool verbatim;
 };
 
+struct shared_umem_len_ctx {
+	u32 short_len;
+	u32 long_len;
+};
+
+struct shared_umem_uneven_dist_ctx {
+	u32 total_pkts;
+	u32 pkt_len;
+};
+
 static inline bool pkt_continues(u32 options)
 {
 	return options & XDP_PKT_CONTD;
@@ -268,6 +285,10 @@ int testapp_xdp_metadata(struct test_spec *test);
 int testapp_xdp_metadata_mb(struct test_spec *test);
 int testapp_xdp_prog_cleanup(struct test_spec *test);
 int testapp_xdp_shared_umem(struct test_spec *test);
+int testapp_shared_umem_4_sockets(struct test_spec *test);
+int testapp_shared_umem_length_based(struct test_spec *test);
+int testapp_shared_umem_uneven_dist(struct test_spec *test);
+int testapp_shared_umem_unaligned(struct test_spec *test);
 
 void *worker_testapp_validate_rx(void *arg);
 void *worker_testapp_validate_tx(void *arg);
@@ -291,6 +312,10 @@ static const struct test_spec tests[] = {
 	{.name = "XDP_PROG_CLEANUP", .test_func = testapp_xdp_prog_cleanup},
 	{.name = "XDP_DROP_HALF", .test_func = testapp_xdp_drop},
 	{.name = "XDP_SHARED_UMEM", .test_func = testapp_xdp_shared_umem},
+	{.name = "SHARED_UMEM_4_SOCKETS", .test_func = testapp_shared_umem_4_sockets},
+	{.name = "SHARED_UMEM_LENGTH_BASED", .test_func = testapp_shared_umem_length_based},
+	{.name = "SHARED_UMEM_UNEVEN_DIST", .test_func = testapp_shared_umem_uneven_dist},
+	{.name = "SHARED_UMEM_UNALIGNED", .test_func = testapp_shared_umem_unaligned},
 	{.name = "XDP_METADATA_COPY", .test_func = testapp_xdp_metadata},
 	{.name = "XDP_METADATA_COPY_MULTI_BUFF", .test_func = testapp_xdp_metadata_mb},
 	{.name = "ALIGNED_INV_DESC_MULTI_BUFF", .test_func = testapp_aligned_inv_desc_mb},
diff --git a/tools/testing/selftests/bpf/progs/xsk_xdp_progs.c b/tools/testing/selftests/bpf/progs/xsk_xdp_progs.c
index 023d8befd4cab..8bab2bcf52ab2 100644
--- a/tools/testing/selftests/bpf/progs/xsk_xdp_progs.c
+++ b/tools/testing/selftests/bpf/progs/xsk_xdp_progs.c
@@ -10,7 +10,7 @@
 
 struct {
 	__uint(type, BPF_MAP_TYPE_XSKMAP);
-	__uint(max_entries, 2);
+	__uint(max_entries, MAX_SOCKETS);
 	__uint(key_size, sizeof(int));
 	__uint(value_size, sizeof(int));
 } xsk SEC(".maps");
@@ -75,6 +75,19 @@ SEC("xdp") int xsk_xdp_shared_umem(struct xdp_md *xdp)
 	return bpf_redirect_map(&xsk, idx, XDP_DROP);
 }
 
+SEC("xdp") int xsk_xdp_shared_umem_length_based(struct xdp_md *xdp)
+{
+	void *data = (void *)(long)xdp->data;
+	void *data_end = (void *)(long)xdp->data_end;
+	__u32 pkt_len = data_end - data;
+
+	/* Route packets by total XDP-visible packet length. */
+	if (pkt_len <= SHARED_UMEM_LEN_SPLIT)
+		return bpf_redirect_map(&xsk, 0, XDP_DROP);
+	else
+		return bpf_redirect_map(&xsk, 1, XDP_DROP);
+}
+
 SEC("xdp.frags") int xsk_xdp_adjust_tail(struct xdp_md *xdp)
 {
 	__u32 buff_len, curr_buff_len;
diff --git a/tools/testing/selftests/bpf/xsk_xdp_common.h b/tools/testing/selftests/bpf/xsk_xdp_common.h
index 45810ff552dae..bd9b5dcf8c8ee 100644
--- a/tools/testing/selftests/bpf/xsk_xdp_common.h
+++ b/tools/testing/selftests/bpf/xsk_xdp_common.h
@@ -3,7 +3,8 @@
 #ifndef XSK_XDP_COMMON_H_
 #define XSK_XDP_COMMON_H_
 
-#define MAX_SOCKETS 2
+#define MAX_SOCKETS 4
+#define SHARED_UMEM_LEN_SPLIT 64
 #define PKT_HDR_ALIGN (sizeof(struct ethhdr) + 2) /* Just to align the data in the packet */
 
 struct xdp_info {
diff --git a/tools/testing/selftests/bpf/xskxceiver.c b/tools/testing/selftests/bpf/xskxceiver.c
index 7dad8556a722e..5d8a29e7224e2 100644
--- a/tools/testing/selftests/bpf/xskxceiver.c
+++ b/tools/testing/selftests/bpf/xskxceiver.c
@@ -117,12 +117,12 @@ static void __exit_with_error(int error, const char *file, const char *func, int
 
 #define exit_with_error(error) __exit_with_error(error, __FILE__, __func__, __LINE__)
 
-static bool ifobj_zc_avail(struct ifobject *ifobject)
+static bool ifobj_zc_avail(struct ifobject *ifobj)
 {
 	size_t umem_sz = DEFAULT_UMEM_BUFFERS * XSK_UMEM__DEFAULT_FRAME_SIZE;
 	int mmap_flags = MAP_PRIVATE | MAP_ANONYMOUS | MAP_NORESERVE;
-	struct xsk_socket_info *xsk;
-	struct xsk_umem_info *umem;
+	struct xsk_socket_info xsk = {};
+	struct xsk_umem_info umem = {};
 	bool zc_avail = false;
 	void *bufs;
 	int ret;
@@ -131,32 +131,26 @@ static bool ifobj_zc_avail(struct ifobject *ifobject)
 	if (bufs == MAP_FAILED)
 		exit_with_error(errno);
 
-	umem = calloc(1, sizeof(struct xsk_umem_info));
-	if (!umem) {
-		munmap(bufs, umem_sz);
-		exit_with_error(ENOMEM);
-	}
-	umem->frame_size = XSK_UMEM__DEFAULT_FRAME_SIZE;
-	ret = xsk_configure_umem(ifobject, umem, bufs, umem_sz);
+	umem.mmap_size = umem_sz;
+	umem.frame_size = XSK_UMEM__DEFAULT_FRAME_SIZE;
+	ret = xsk_configure_umem(ifobj, &umem, bufs, umem_sz);
 	if (ret)
 		exit_with_error(-ret);
 
-	xsk = calloc(1, sizeof(struct xsk_socket_info));
-	if (!xsk)
-		goto out;
-	ifobject->bind_flags = XDP_USE_NEED_WAKEUP | XDP_ZEROCOPY;
-	ifobject->rx_on = true;
-	xsk->rxqsize = XSK_RING_CONS__DEFAULT_NUM_DESCS;
-	ret = xsk_configure_socket(xsk, umem, ifobject, false);
+	xsk.umem_ref = true;
+
+	ifobj->bind_flags = XDP_USE_NEED_WAKEUP | XDP_ZEROCOPY;
+	ifobj->rx_on = true;
+	xsk.rxqsize = XSK_RING_CONS__DEFAULT_NUM_DESCS;
+	ret = xsk_configure_socket(&xsk, &umem, ifobj, false);
 	if (!ret)
 		zc_avail = true;
 
-	xsk_socket__delete(xsk->xsk);
-	free(xsk);
-out:
-	munmap(umem->buffer, umem_sz);
-	xsk_umem__delete(umem->umem);
-	free(umem);
+	/* Use the same refcount-based teardown path for both success and failure. */
+	ret = xsk_delete_socket(&xsk);
+	if (ret)
+		exit_with_error(-ret);
+
 	return zc_avail;
 }