Merge pull request iovisor#867 from shodoco/extern

Don't close extern table when a module destructs

src/cc/bpf_module.cc

@@ -118,10 +118,11 @@ 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);
+ }
  }
  }
 }

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,15 +678,15 @@ 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;
  }
  table_it->is_shared = true;
  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,21 +36,28 @@ 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;
 }
 
 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 @@ def cleanup(self):
  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

@@ -0,0 +1,23 @@
+#!/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()