From 0b69111376a6ac668548735e3c095de06af9a134 Mon Sep 17 00:00:00 2001
From: humaite <humaite@arweave>
Date: Fri, 4 Oct 2024 16:39:06 +0200
Subject: [PATCH] add tests and use ets:select instead of comprehension list

---
 apps/arweave/src/ar_peers.erl | 73 ++++++++++++++++++++++++++---------
 1 file changed, 55 insertions(+), 18 deletions(-)

diff --git a/apps/arweave/src/ar_peers.erl b/apps/arweave/src/ar_peers.erl
index 2742a4494..76ebd9703 100644
--- a/apps/arweave/src/ar_peers.erl
+++ b/apps/arweave/src/ar_peers.erl
@@ -175,20 +175,26 @@ start_link() ->
 %% @end
 %%--------------------------------------------------------------------
 get_peers(connected) ->
-	Peers = get_peers(lifetime),
-	[ P || P <- Peers, is_connected_peer(P) ];
+	Tag = {connection, active},
+	Pattern = {{ar_tags, ?MODULE, '$1', Tag}, '$3'},
+	Guard = [{'=:=', '$3', true}],
+	Select = ['$1'],
+	ets:select(?MODULE, [{Pattern, Guard, Select}]);
+get_peers({timestamp, Seconds}) when is_integer(Seconds) ->
+	Timefilter = erlang:system_time(seconds) - Seconds,
+	Tag = {connection, last},
+	Pattern = {{ar_tags, ?MODULE, '$1', Tag}, '$3'},
+	Guard = [{'>=', '$3', Timefilter}],
+	Select = ['$1'],
+	ets:select(?MODULE, [{Pattern, Guard, Select}]);
 get_peers({hours, Hour}) ->
-	Peers = get_peers(lifetime),
-	[ P || P <- Peers, get_connection_timestamp_peer(P) >= Hour*60*60 ];
+	get_peers({timestamp, Hour*60*60});
 get_peers({days, Day}) ->
-	Peers = get_peers(lifetime),
-	[ P || P <- Peers, get_connection_timestamp_peer(P) >= Day*86400 ];
+	get_peers({timestamp, Day*86400});
 get_peers({weeks, Week}) ->
-	Peers = get_peers(lifetime),
-	[ P || P <- Peers, get_connection_timestamp_peer(P) >= Week*86400*7 ];
+	get_peers({timestamp, Week*86400*7});
 get_peers({months, Month}) ->
-	Peers = get_peers(lifetime),
-	[ P || P <- Peers, get_connection_timestamp_peer(P) >= Month*86400*7*30 ];
+	get_peers({timestamp, Month*86400*7*30});
 get_peers(Ranking) ->
 	case catch ets:lookup(?MODULE, {peers, Ranking}) of
 		{'EXIT', _} ->
@@ -1016,33 +1022,64 @@ is_connected_peer(Peer) ->
 %%% Tests.
 %%%===================================================================
 connected_peer_test() ->
-	ets:delete_all_objects(?MODULE),
 	Peer = {100, 117, 109, 98, 1234},
 
-	% get_peers(active) returns an empty list when no
-	% active nodes are present
+	% drop all objects from the table to start with a clean state
+	ets:delete_all_objects(?MODULE),
+
+
+	% get all peers connected, it should returns nothing by
+	% default because the table is empty.
 	?assertEqual([], get_peers(connected)),
 	?assertEqual(undefined, get_connection_timestamp_peer(Peer)),
 
-	% by default, if a peer is not present in the table, it must
-	% be down.
+	% manually add a new peer using set_ranked_peers function.
+	% the node is not connected because gun did not manage the
+	% connection in this test.
 	set_ranked_peers(lifetime, [Peer]),
 	?assertEqual(false, is_connected_peer(Peer)),
 	?assertEqual(undefined, get_connection_timestamp_peer(Peer)),
 
-	% a peer is connected
+	% force this peer to be connected using connected_peer
+	% function. A timestamp is created.
 	connected_peer(Peer),
 	Timestamp = get_connection_timestamp_peer(Peer),
 	?assertEqual(true, is_connected_peer(Peer)),
 	?assertEqual(Timestamp, get_connection_timestamp_peer(Peer)),
 	?assertNotEqual(undefined, get_connection_timestamp_peer(Peer)),
 	?assertEqual([Peer], get_peers(connected)),
+	?assertEqual([Peer], get_peers({hours, 2})),
+	?assertEqual([Peer], get_peers({days, 3})),
+	?assertEqual([Peer], get_peers({weeks, 4})),
+	?assertEqual([Peer], get_peers({months, 5})),
 
-	% a peer is disconnected
+	% Now remove the connection to the peer. A timestamp must
+	% still be there.
 	disconnected_peer(Peer),
 	?assertEqual(false, is_connected_peer(Peer)),
 	?assertNotEqual(undefined, get_connection_timestamp_peer(Peer)),
-	?assertEqual([], get_peers(connected)).
+	?assertEqual([], get_peers(connected)),
+	?assertEqual([Peer], get_peers({hours, 2})),
+	?assertEqual([Peer], get_peers({days, 3})),
+	?assertEqual([Peer], get_peers({weeks, 4})),
+	?assertEqual([Peer], get_peers({months, 5})),
+
+	% let modify manually the timestamp to check get_peers/1
+	% function, and overwrite Peer timestamp with some defined
+	% values.
+	Time = erlang:system_time(second),
+	% set its tag to 11 hour
+	set_tag(Peer, {connection, last}, Time-(11*60*60)),
+	?assertEqual([], get_peers({hours, 2})),
+	% set its tag to 11 days
+	set_tag(Peer, {connection, last}, Time-(11*86400)),
+	?assertEqual([], get_peers({days, 2})),
+	% set its tag to 11 weeks
+	set_tag(Peer, {connection, last}, Time-(11*86400*7)),
+	?assertEqual([], get_peers({weeks, 2})),
+	% set its tag to 11 months
+	set_tag(Peer, {connection, last}, Time-(11*86400*30)),
+	?assertEqual([], get_peers({months, 2})).
 
 rotate_peer_ports_test() ->
 	Peer = {2, 2, 2, 2, 1},