dynamically set initial filter map size

Yelp · Mar 11, 2022 · 109f25a · 109f25a
1 parent 34dc02f
commit 109f25a
Show file tree

Hide file tree

Showing 6 changed files with 26 additions and 3 deletions.
diff --git a/pidtree_bcc/probes/__init__.py b/pidtree_bcc/probes/__init__.py
@@ -22,6 +22,7 @@
 from pidtree_bcc.plugins import load_plugins
 from pidtree_bcc.utils import find_subclass
 from pidtree_bcc.utils import never_crash
+from pidtree_bcc.utils import round_nearest_multiple
 
 
 class BPFProbe:
@@ -45,6 +46,8 @@ class BPFProbe:
     # (not via Jinja-templated if statements)
     USES_DYNAMIC_FILTERS = False
     NET_FILTER_MAP_NAME = 'net_filter_map'
+    NET_FILTER_MAP_SIZE_MAX = 4 * 1024
+    NET_FILTER_MAP_SIZE_SCALING = 512
     PORT_FILTER_MAP_NAME = 'port_filter_map'
 
     def __init__(
@@ -119,6 +122,10 @@ def build_probe_config(self, probe_config: dict, hotswap_only: bool = False) ->
                 template_config['NET_FILTER_MAP_NAME'] = self.NET_FILTER_MAP_NAME
                 template_config['PORT_FILTER_MAP_NAME'] = self.PORT_FILTER_MAP_NAME
                 template_config['NET_FILTER_MAX_PORT_RANGES'] = NET_FILTER_MAX_PORT_RANGES
+                template_config['NET_FILTER_MAP_SIZE'] = min(
+                    self.NET_FILTER_MAP_SIZE_MAX,
+                    round_nearest_multiple(len(self.net_filters), self.NET_FILTER_MAP_SIZE_SCALING, headroom=128),
+                )
         return template_config
 
     def _process_events(self, cpu: Any, data: Any, size: Any, from_bpf: bool = True):

diff --git a/pidtree_bcc/probes/net_listen.j2 b/pidtree_bcc/probes/net_listen.j2
@@ -12,7 +12,7 @@ struct listen_bind_t {
     u8  protocol;
 };
 
-{{ utils.net_filter_trie_init(NET_FILTER_MAP_NAME, PORT_FILTER_MAP_NAME, max_ports=NET_FILTER_MAX_PORT_RANGES) }}
+{{ utils.net_filter_trie_init(NET_FILTER_MAP_NAME, PORT_FILTER_MAP_NAME, size=NET_FILTER_MAP_SIZE, max_ports=NET_FILTER_MAX_PORT_RANGES) }}
 
 {{ utils.get_proto_func() }}
 

diff --git a/pidtree_bcc/probes/tcp_connect.j2 b/pidtree_bcc/probes/tcp_connect.j2
@@ -12,7 +12,7 @@ struct connection_t {
     u16 dport;
 };
 
-{{ utils.net_filter_trie_init(NET_FILTER_MAP_NAME, PORT_FILTER_MAP_NAME, max_ports=NET_FILTER_MAX_PORT_RANGES) }}
+{{ utils.net_filter_trie_init(NET_FILTER_MAP_NAME, PORT_FILTER_MAP_NAME, size=NET_FILTER_MAP_SIZE, max_ports=NET_FILTER_MAX_PORT_RANGES) }}
 
 int kprobe__tcp_v4_connect(struct pt_regs *ctx, struct sock *sk)
 {

diff --git a/pidtree_bcc/probes/udp_session.j2 b/pidtree_bcc/probes/udp_session.j2
@@ -17,7 +17,7 @@ struct udp_session_event {
 BPF_PERF_OUTPUT(events);
 BPF_HASH(tracing, u64, u8);
 
-{{ utils.net_filter_trie_init(NET_FILTER_MAP_NAME, PORT_FILTER_MAP_NAME, max_ports=NET_FILTER_MAX_PORT_RANGES) }}
+{{ utils.net_filter_trie_init(NET_FILTER_MAP_NAME, PORT_FILTER_MAP_NAME, size=NET_FILTER_MAP_SIZE, max_ports=NET_FILTER_MAX_PORT_RANGES) }}
 
 {{ utils.get_proto_func() }}
 

diff --git a/pidtree_bcc/utils.py b/pidtree_bcc/utils.py
@@ -144,3 +144,14 @@ def get_network_namespace(pid: int = None) -> int:
 def self_restart():
     """ Causes pidtree-bcc to restart itself """
     os.kill(os.getpid(), signal.SIGHUP)
+
+
+def round_nearest_multiple(value: int, factor: int, headroom: int = 0) -> int:
+    """ Round value to nearest multiple given a factor
+
+    :param int value: starting value
+    :param int factor: factor to use in finding multiple
+    :param int headroom: ensure this much difference between value and result
+    :return: rounded value
+    """
+    return factor * ((value + headroom) // factor + 1)
diff --git a/tests/utils_test.py b/tests/utils_test.py
@@ -52,3 +52,8 @@ def test_netmask_to_prefixlen():
     assert utils.netmask_to_prefixlen('255.0.0.0') == 8
     with pytest.raises(ValueError):
         utils.netmask_to_prefixlen('1.1.1.1')
+
+
+def test_round_nearest_multiple():
+    assert utils.round_nearest_multiple(23, 10) == 30
+    assert utils.round_nearest_multiple(27, 10, 4) == 40