add table type to shared tables

src/cc/frontends/clang/b_frontend_action.cc

@@ -668,6 +668,7 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
     } else if (A->getName() == "maps/extern") {
       is_extern = true;
       table.fd = SharedTables::instance()->lookup_fd(table.name);
+      table.type = SharedTables::instance()->lookup_type(table.name);
     } else if (A->getName() == "maps/export") {
       if (table.name.substr(0, 2) == "__")
         table.name = table.name.substr(2);
@@ -678,7 +679,7 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
         error(Decl->getLocStart(), "reference to undefined table");
         return false;
       }
-      if (!SharedTables::instance()->insert_fd(table.name, table_it->fd)) {
+      if (!SharedTables::instance()->insert_fd(table.name, table_it->fd, table_it->type)) {
         error(Decl->getLocStart(), "could not export bpf map %0: %1") << table.name << "already in use";
         return false;
       }

src/cc/shared_table.cc

@@ -17,6 +17,7 @@
 #include <unistd.h>
 
 #include "shared_table.h"
+#include "compat/linux/bpf.h"
 
 namespace ebpf {
 
@@ -35,13 +36,20 @@ int SharedTables::lookup_fd(const string &name) const {
   auto table = tables_.find(name);
   if (table == tables_.end())
     return -1;
-  return table->second;
+  return table->second.first;
 }
 
-bool SharedTables::insert_fd(const string &name, int fd) {
+int SharedTables::lookup_type(const string &name) const {
+  auto table = tables_.find(name);
+  if (table == tables_.end())
+    return BPF_MAP_TYPE_UNSPEC;
+  return table->second.second;
+}
+
+bool SharedTables::insert_fd(const string &name, int fd, int type) {
   if (tables_.find(name) != tables_.end())
     return false;
-  tables_[name] = fd;
+  tables_[name] = std::make_pair(fd, type);
   return true;
 }
 
@@ -49,7 +57,7 @@ bool SharedTables::remove_fd(const string &name) {
   auto table = tables_.find(name);
   if (table == tables_.end())
     return false;
-  close(table->second);
+  close(table->second.first);
   tables_.erase(table);
   return true;
 }

src/cc/shared_table.h

@@ -27,14 +27,16 @@ class SharedTables {
  public:
   static SharedTables * instance();
   // add an fd to the shared table, return true if successfully inserted
-  bool insert_fd(const std::string &name, int fd);
+  bool insert_fd(const std::string &name, int fd, int type);
   // lookup an fd in the shared table, or -1 if not found
   int lookup_fd(const std::string &name) const;
+  // lookup on map type in the shared table, or BPF_MAP_TYPE_UNSPEC if not found
+  int lookup_type(const std::string &name) const;
   // close and remove a shared fd. return true if the value was found
   bool remove_fd(const std::string &name);
  private:
   static SharedTables *instance_;
-  std::map<std::string, int> tables_;
+  std::map<std::string, std::pair<int, int>> tables_;
 };
 
 }

src/python/bcc/__init__.py

@@ -1058,5 +1058,11 @@ class BPF(object):
             lib.bpf_module_destroy(self.module)
             self.module = None
 
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.cleanup()
+
 
 from .usdt import USDT

tests/python/test_shared_table.py

+#!/usr/bin/env python
+# Copyright (c) PLUMgrid, Inc.
+# Licensed under the Apache License, Version 2.0 (the "License")
+
+import ctypes as ct
+import unittest
+from bcc import BPF
+
+class TestSharedTable(unittest.TestCase):
+    def test_close_extern(self):
+        b1 = BPF(text="""BPF_TABLE_PUBLIC("array", int, int, table1, 10);""")
+
+        with BPF(text="""BPF_TABLE("extern", int, int, table1, 10);""") as b2:
+            t2 = b2["table1"]
+            t2[ct.c_int(1)] = ct.c_int(10)
+            self.assertEqual(len(t2), 10)
+
+        t1 = b1["table1"]
+        self.assertEqual(t1[ct.c_int(1)].value, 10)
+        self.assertEqual(len(t1), 10)
+
+if __name__ == "__main__":
+    unittest.main()