Repository: cilium/ebpf
Pull request #47 (2020-01-27 11:19:59)

collection: allow rewriting individual maps
Given a CollectionSpec, it's sometimes useful to use existing Maps instead
of having NewCollection spec create them from scratch. For example,
multiple programs might share the same map. Another use case is persistet
state that needs to be kept between two versions of an eBPF program.

Add a helper RewriteMaps which allows to achieve this.

lmb: cc @hamburgerbun

Patch:
diff --git a/collection.go b/collection.go
index bf6a96d..d96b670 100644
--- a/collection.go
+++ b/collection.go
@@ -39,6 +39,43 @@ func (cs *CollectionSpec) Copy() *CollectionSpec {
 	return &cpy
 }
 
+// RewriteMaps replaces all references to specific maps.
+//
+// Use this function to use pre-existing maps instead of creating new ones
+// when calling NewCollection. Any named maps are removed from CollectionSpec.Maps.
+//
+// Returns an error if a named map isn't used in at least one program.
+func (cs *CollectionSpec) RewriteMaps(maps map[string]*Map) error {
+	for symbol, m := range maps {
+		// have we seen a program that uses this symbol / map
+		seen := false
+		fd := m.FD()
+		for progName, progSpec := range cs.Programs {
+			err := progSpec.Instructions.RewriteMapPtr(symbol, fd)
+
+			switch {
+			case err == nil:
+				seen = true
+
+			case asm.IsUnreferencedSymbol(err):
+				// Not all programs need to use the map
+
+			default:
+				return errors.Wrapf(err, "program %s", progName)
+			}
+		}
+
+		if !seen {
+			return errors.Errorf("map %s not referenced by any programs", symbol)
+		}
+
+		// Prevent NewCollection from creating rewritten maps
+		delete(cs.Maps, symbol)
+	}
+
+	return nil
+}
+
 // Collection is a collection of Programs and Maps associated
 // with their symbols
 type Collection struct {
diff --git a/map.go b/map.go
index 86eff1a..586cfff 100644
--- a/map.go
+++ b/map.go
@@ -129,11 +129,11 @@ func createMap(spec *MapSpec, inner *internal.FD, handle *btf.Handle) (*Map, err
 		spec.ValueSize = 4
 
 	case PerfEventArray:
-		if spec.KeySize != 0 {
-			return nil, errors.Errorf("KeySize must be zero for perf event array")
+		if spec.KeySize != 0 && spec.KeySize != 4 {
+			return nil, errors.Errorf("KeySize must be zero or four for perf event array")
 		}
-		if spec.ValueSize != 0 {
-			return nil, errors.Errorf("ValueSize must be zero for perf event array")
+		if spec.ValueSize != 0 && spec.ValueSize != 4 {
+			return nil, errors.Errorf("ValueSize must be zero or four for perf event array")
 		}
 		if spec.MaxEntries == 0 {
 			n, err := internal.OnlineCPUs()


Test patch:
diff --git a/collection_test.go b/collection_test.go
index 1015c32..e1b9a17 100644
--- a/collection_test.go
+++ b/collection_test.go
@@ -80,7 +80,7 @@ func TestCollectionSpecCopy(t *testing.T) {
 	}
 }
 
-func TestCollectionSpecOverwriteMaps(t *testing.T) {
+func TestCollectionSpecRewriteMaps(t *testing.T) {
 	insns := asm.Instructions{
 		// R1 map
 		asm.LoadMapPtr(asm.R1, 0),
@@ -126,22 +126,22 @@ func TestCollectionSpecOverwriteMaps(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	err = cs.Programs["test-prog"].Instructions.RewriteMapPtr("test-map", newMap.FD())
+	err = cs.RewriteMaps(map[string]*Map{
+		"test-map": newMap,
+	})
 	if err != nil {
 		t.Fatal(err)
 	}
 
-	coll, err := NewCollection(cs)
-	if err != nil {
-		t.Fatal(err)
+	if cs.Maps["test-map"] != nil {
+		t.Error("RewriteMaps doesn't remove map from CollectionSpec.Maps")
 	}
-	defer coll.Close()
 
-	oldMap := coll.Maps["test-map"]
-	err = oldMap.Put(uint32(0), uint32(5))
+	coll, err := NewCollection(cs)
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer coll.Close()
 
 	ret, _, err := coll.Programs["test-prog"].Test(make([]byte, 14))
 	testutils.SkipIfNotSupported(t, err)
@@ -149,8 +149,6 @@ func TestCollectionSpecOverwriteMaps(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	t.Log(ret)
-
 	if ret != 2 {
 		t.Fatal("new / override map not used")
 	}
diff --git a/map_test.go b/map_test.go
index e71dd1c..0b86349 100644
--- a/map_test.go
+++ b/map_test.go
@@ -292,13 +292,20 @@ func TestNewMapInMapFromFD(t *testing.T) {
 }
 
 func TestPerfEventArray(t *testing.T) {
-	m, err := NewMap(&MapSpec{
-		Type: PerfEventArray,
-	})
-	if err != nil {
-		t.Fatal("Can't create perf event array:", err)
+	specs := []*MapSpec{
+		{Type: PerfEventArray},
+		{Type: PerfEventArray, KeySize: 4},
+		{Type: PerfEventArray, ValueSize: 4},
+	}
+
+	for _, spec := range specs {
+		m, err := NewMap(spec)
+		if err != nil {
+			t.Errorf("Can't create perf event array from %v: %s", spec, err)
+		} else {
+			m.Close()
+		}
 	}
-	m.Close()
 }
 
 func createMapInMap(t *testing.T, typ MapType) *Map {