Merge pull request #393 from aztecher/impl-set_max_entries

aya: add BpfLoader::set_max_entries

aya/src/bpf.rs

@@ -191,6 +191,7 @@ pub struct BpfLoader<'a> {
     btf: Option<Cow<'a, Btf>>,
     map_pin_path: Option<PathBuf>,
     globals: HashMap<&'a str, &'a [u8]>,
+    max_entries: HashMap<&'a str, u32>,
     features: Features,
     extensions: HashSet<&'a str>,
     verifier_log_level: VerifierLogLevel,
@@ -227,6 +228,7 @@ impl<'a> BpfLoader<'a> {
             btf: Btf::from_sys_fs().ok().map(Cow::Owned),
             map_pin_path: None,
             globals: HashMap::new(),
+            max_entries: HashMap::new(),
             features,
             extensions: HashSet::new(),
             verifier_log_level: VerifierLogLevel::default(),
@@ -315,6 +317,27 @@ impl<'a> BpfLoader<'a> {
         self
     }
 
+    /// Set the max_entries for specified map.
+    ///
+    /// Overwrite the value of max_entries of the map that matches
+    /// the provided name before the map is created.
+    ///
+    /// # Example
+    ///
+    /// ```no_run
+    /// use aya::BpfLoader;
+    ///
+    /// let bpf = BpfLoader::new()
+    ///     .set_max_entries("map", 64)
+    ///     .load_file("file.o")?;
+    /// # Ok::<(), aya::BpfError>(())
+    /// ```
+    ///
+    pub fn set_max_entries(&mut self, name: &'a str, size: u32) -> &mut BpfLoader<'a> {
+        self.max_entries.insert(name, size);
+        self
+    }
+
     /// Treat the provided program as an [`Extension`]
     ///
     /// When attempting to load the program with the provided `name`
@@ -409,10 +432,14 @@ impl<'a> BpfLoader<'a> {
         if let Some(btf) = &self.btf {
             obj.relocate_btf(btf)?;
         }
-
         let mut maps = HashMap::new();
         for (name, mut obj) in obj.maps.drain() {
-            if obj.map_type() == BPF_MAP_TYPE_PERF_EVENT_ARRAY as u32 && obj.max_entries() == 0 {
+            match self.max_entries.get(name.as_str()) {
+                Some(size) => obj.set_max_entries(*size),
+                None => {
+                    if obj.map_type() == BPF_MAP_TYPE_PERF_EVENT_ARRAY as u32
+                        && obj.max_entries() == 0
+                    {
                         obj.set_max_entries(
                             possible_cpus()
                                 .map_err(|error| BpfError::FileError {
@@ -422,6 +449,8 @@ impl<'a> BpfLoader<'a> {
                                 .len() as u32,
                         );
                     }
+                }
+            }
             let mut map = Map {
                 obj,
                 fd: None,