Merge pull request #867 from shodoco/extern

Don't close extern table when a module destructs

src/cc/bpf_module.cc

@@ -118,12 +118,13 @@ BPFModule::~BPFModule() {
   ctx_.reset();
   if (tables_) {
     for (auto table : *tables_) {
-      if (table.is_shared)
+      if (table.is_shared) {
         SharedTables::instance()->remove_fd(table.name);
-      else
+      } else if (!table.is_extern) {
         close(table.fd);
       }
     }
+  }
 }
 
 static void debug_printf(Module *mod, IRBuilder<> &B, const string &fmt, vector<Value *> args) {

src/cc/frontends/clang/b_frontend_action.cc

@@ -631,7 +631,6 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       ++i;
     }
 
-    bool is_extern = false;
     bpf_map_type map_type = BPF_MAP_TYPE_UNSPEC;
     if (A->getName() == "maps/hash") {
       map_type = BPF_MAP_TYPE_HASH;
@@ -666,8 +665,9 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
     } else if (A->getName() == "maps/stacktrace") {
       map_type = BPF_MAP_TYPE_STACK_TRACE;
     } else if (A->getName() == "maps/extern") {
-      is_extern = true;
+      table.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 +678,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;
       }
@@ -686,7 +686,7 @@ bool BTypeVisitor::VisitVarDecl(VarDecl *Decl) {
       return true;
     }
 
-    if (!is_extern) {
+    if (!table.is_extern) {
       if (map_type == BPF_MAP_TYPE_UNSPEC) {
         error(Decl->getLocStart(), "unsupported map type: %0") << A->getName();
         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/cc/table_desc.h

@@ -40,6 +40,7 @@ struct TableDesc {
   llvm::Function *key_snprintf;
   llvm::Function *leaf_snprintf;
   bool is_shared;
+  bool is_extern;
 };
 
 }  // namespace ebpf

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) 2016 Facebook, 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()