Fixes #2518 -- weird behaviour of lookup_or_init (#2520)

* Allow lookup_or_init to return NULL (rather than just returning from the current function)
* Fixed a couple of bad edits found when running tests. Also fixed a bug in the
test runner. Also fixed a bug in libbcc where the python signature did not match
the actual implementation.

Author: pjsg

docs/reference_guide.md

@@ -767,7 +767,7 @@ Examples in situ:
 
 Syntax: ```*val map.lookup_or_init(&key, &zero)```
 
-Lookup the key in the map, and return a pointer to its value if it exists, else initialize the key's value to the second argument. This is often used to initialize values to zero.
+Lookup the key in the map, and return a pointer to its value if it exists, else initialize the key's value to the second argument. This is often used to initialize values to zero. If the key cannot be inserted (e.g. the map is full) then NULL is returned.
 
 Examples in situ:
 [search /examples](https://github.com/iovisor/bcc/search?q=lookup_or_init+path%3Aexamples&type=Code),

docs/tutorial_bcc_python_developer.md

@@ -559,7 +559,9 @@ int count(struct pt_regs *ctx) {
     bpf_probe_read(&key.c, sizeof(key.c), (void *)PT_REGS_PARM1(ctx));
     // could also use `counts.increment(key)`
     val = counts.lookup_or_init(&key, &zero);
-    (*val)++;
+    if (val) {
+      (*val)++;
+    }
     return 0;
 };
 """)
@@ -678,7 +680,9 @@ int count_sched(struct pt_regs *ctx, struct task_struct *prev) {
 
     // could also use `stats.increment(key);`
     val = stats.lookup_or_init(&key, &zero);
-    (*val)++;
+    if (val) {
+      (*val)++;
+    }
     return 0;
 }
 ```

examples/cpp/LLCStat.cc

@@ -43,7 +43,9 @@ int on_cache_miss(struct bpf_perf_event_data *ctx) {
 
     u64 zero = 0, *val;
     val = miss_count.lookup_or_init(&key, &zero);
-    (*val) += ctx->sample_period;
+    if (val) {
+        (*val) += ctx->sample_period;
+    }
 
     return 0;
 }
@@ -54,7 +56,9 @@ int on_cache_ref(struct bpf_perf_event_data *ctx) {
 
     u64 zero = 0, *val;
     val = ref_count.lookup_or_init(&key, &zero);
-    (*val) += ctx->sample_period;
+    if (val) {
+        (*val) += ctx->sample_period;
+    }
 
     return 0;
 }

examples/cpp/TCPSendStack.cc

@@ -39,7 +39,9 @@ int on_tcp_send(struct pt_regs *ctx) {
 
   u64 zero = 0, *val;
   val = counts.lookup_or_init(&key, &zero);
-  (*val)++;
+  if (val) {
+    (*val)++;
+  }
 
   return 0;
 }

examples/cpp/UseExternalMap.cc

@@ -56,7 +56,9 @@ int on_sched_switch(struct tracepoint__sched__sched_switch *args) {
   __builtin_memcpy(&key.prev_comm, args->prev_comm, 16);
   __builtin_memcpy(&key.next_comm, args->next_comm, 16);
   val = counts.lookup_or_init(&key, &zero);
-  (*val)++;
+  if (val) {
+    (*val)++;
+  }
 
   return 0;
 }

examples/lua/offcputime.lua

@@ -65,7 +65,9 @@ int oncpu(struct pt_regs *ctx, struct task_struct *prev) {
     key.stack_id = stack_traces.get_stackid(ctx, stack_flags);
 
     val = counts.lookup_or_init(&key, &zero);
-    (*val) += delta;
+    if (val) {
+        (*val) += delta;
+    }
     return 0;
 }
 ]]

examples/lua/task_switch.lua

@@ -33,7 +33,9 @@ int count_sched(struct pt_regs *ctx, struct task_struct *prev) {
   key.prev_pid = prev->pid;
 
   val = stats.lookup_or_init(&key, &zero);
-  (*val)++;
+  if (val) {
+    (*val)++;
+  }
   return 0;
 }
 ]]

examples/networking/distributed_bridge/tunnel.c

@@ -38,7 +38,9 @@ int handle_ingress(struct __sk_buff *skb) {
     struct vni_key vk = {ethernet->src, *ifindex, 0};
     struct host *src_host = mac2host.lookup_or_init(&vk,
         &(struct host){tkey.tunnel_id, tkey.remote_ipv4, 0, 0});
-    lock_xadd(&src_host->rx_pkts, 1);
+    if (src_host) {
+      lock_xadd(&src_host->rx_pkts, 1);
+    }
     bpf_clone_redirect(skb, *ifindex, 1/*ingress*/);
   } else {
     bpf_trace_printk("ingress invalid tunnel_id=%d\n", tkey.tunnel_id);

examples/networking/tunnel_monitor/monitor.c

@@ -126,12 +126,14 @@ ip: ;
     swap_ipkey(&key);
   struct counters zleaf = {0};
   struct counters *leaf = stats.lookup_or_init(&key, &zleaf);
-  if (is_ingress) {
-    lock_xadd(&leaf->rx_pkts, 1);
-    lock_xadd(&leaf->rx_bytes, skb->len);
-  } else {
-    lock_xadd(&leaf->tx_pkts, 1);
-    lock_xadd(&leaf->tx_bytes, skb->len);
+  if (leaf) {
+    if (is_ingress) {
+      lock_xadd(&leaf->rx_pkts, 1);
+      lock_xadd(&leaf->rx_bytes, skb->len);
+    } else {
+      lock_xadd(&leaf->tx_pkts, 1);
+      lock_xadd(&leaf->tx_bytes, skb->len);
+    }
   }
   return 1;
 }

examples/networking/vlan_learning/vlan_learning.c

@@ -33,10 +33,12 @@ int handle_phys2virt(struct __sk_buff *skb) {
       int out_ifindex = leaf->out_ifindex;
       struct ifindex_leaf_t zleaf = {0};
       struct ifindex_leaf_t *out_leaf = egress.lookup_or_init(&out_ifindex, &zleaf);
-      // to capture potential configuration changes
-      out_leaf->out_ifindex = skb->ifindex;
-      out_leaf->vlan_tci = skb->vlan_tci;
-      out_leaf->vlan_proto = skb->vlan_proto;
+      if (out_leaf) {
+	// to capture potential configuration changes
+	out_leaf->out_ifindex = skb->ifindex;
+	out_leaf->vlan_tci = skb->vlan_tci;
+	out_leaf->vlan_proto = skb->vlan_proto;
+      }
       // pop the vlan header and send to the destination
       bpf_skb_vlan_pop(skb);
       bpf_clone_redirect(skb, leaf->out_ifindex, 0);

examples/tracing/mallocstacks.py

@@ -47,7 +47,9 @@
     // could also use `calls.increment(key, size);`
     u64 zero = 0, *val;
     val = calls.lookup_or_init(&key, &zero);
-    (*val) += size;
+    if (val) {
+      (*val) += size;
+    }
     return 0;
 };
 """)

examples/tracing/strlen_count.py

@@ -34,7 +34,9 @@
     bpf_probe_read(&key.c, sizeof(key.c), (void *)PT_REGS_PARM1(ctx));
     // could also use `counts.increment(key)`
     val = counts.lookup_or_init(&key, &zero);
-    (*val)++;
+    if (val) {
+      (*val)++;
+    }
     return 0;
 };
 """)

examples/tracing/task_switch.c

@@ -16,6 +16,8 @@ int count_sched(struct pt_regs *ctx, struct task_struct *prev) {
 
     // could also use `stats.increment(key);`
     val = stats.lookup_or_init(&key, &zero);
-    (*val)++;
+    if (val) {
+        (*val)++;
+    }
     return 0;
 }

src/cc/frontends/clang/b_frontend_action.cc

@@ -782,7 +782,6 @@ bool BTypeVisitor::VisitCallExpr(CallExpr *Call) {
           txt += "if (!leaf) {";
           txt += " " + update + ", " + arg0 + ", " + arg1 + ", BPF_NOEXIST);";
           txt += " leaf = " + lookup + ", " + arg0 + ");";
-          txt += " if (!leaf) return 0;";
           txt += "}";
           txt += "leaf;})";
         } else if (memb_name == "increment") {

src/python/bcc/libbcc.py

@@ -89,7 +89,7 @@
 lib.bpf_attach_socket.argtypes = [ct.c_int, ct.c_int]
 lib.bcc_func_load.restype = ct.c_int
 lib.bcc_func_load.argtypes = [ct.c_void_p, ct.c_int, ct.c_char_p, ct.c_void_p,
-        ct.c_size_t, ct.c_char_p, ct.c_uint, ct.c_int, ct.c_char_p, ct.c_uint]
+        ct.c_size_t, ct.c_char_p, ct.c_uint, ct.c_int, ct.c_char_p, ct.c_uint, ct.c_char_p]
 _RAW_CB_TYPE = ct.CFUNCTYPE(None, ct.py_object, ct.c_void_p, ct.c_int)
 _LOST_CB_TYPE = ct.CFUNCTYPE(None, ct.py_object, ct.c_ulonglong)
 lib.bpf_attach_kprobe.restype = ct.c_int

tests/cc/test_bpf_table.cc

@@ -182,7 +182,9 @@ TEST_CASE("test bpf stack table", "[bpf_stack_table]") {
       int stack_id = stack_traces.get_stackid(ctx, BPF_F_REUSE_STACKID);
       int zero = 0, *val;
       val = id.lookup_or_init(&zero, &stack_id);
-      (*val) = stack_id;
+      if (val) {
+        (*val) = stack_id;
+      }
 
       return 0;
     }
@@ -233,7 +235,9 @@ TEST_CASE("test bpf stack_id table", "[bpf_stack_table]") {
       int stack_id = stack_traces.get_stackid(ctx, BPF_F_USER_STACK);
       int zero = 0, *val;
       val = id.lookup_or_init(&zero, &stack_id);
-      (*val) = stack_id;
+      if (val) {
+        (*val) = stack_id;
+      }
 
       return 0;
     }

tests/lua/test_clang.lua

@@ -244,7 +244,9 @@ int kprobe__finish_task_switch(struct pt_regs *ctx, struct task_struct *prev) {
   key.prev_pid = prev->pid;
 
   val = stats.lookup_or_init(&key, &zero);
-  (*val)++;
+  if (val) {
+    (*val)++;
+  }
   return 0;
 }
 ]]}

tests/python/test_clang.py

@@ -397,7 +397,9 @@ def test_task_switch(self):
   key.prev_pid = prev->pid;
 
   val = stats.lookup_or_init(&key, &zero);
-  (*val)++;
+  if (val) {
+    (*val)++;
+  }
   return 0;
 }
 """)

tests/python/test_license.py

@@ -43,7 +43,9 @@ class TestLicense(unittest.TestCase):
     bpf_get_current_comm(&(key.comm), 16);
 
     val = counts.lookup_or_init(&key, &zero);  // update counter
-    (*val)++;
+    if (val) {
+        (*val)++;
+    }
     return 0;
 }
 """

tests/python/test_lru.py

@@ -27,7 +27,9 @@ def test_lru_percpu_hash(self):
             u32 key=0;
             u32 value = 0, *val;
             val = stats.lookup_or_init(&key, &value);
-            *val += 1;
+            if (val) {
+                *val += 1;
+            }
             return 0;
         }
         """

tests/python/test_percpu.py

@@ -30,7 +30,9 @@ def test_u64(self):
             u32 key=0;
             u64 value = 0, *val;
             val = stats.lookup_or_init(&key, &value);
-            *val += 1;
+            if (val) {
+                *val += 1;
+            }
             return 0;
         }
         """
@@ -60,7 +62,9 @@ def test_u32(self):
             u32 key=0;
             u32 value = 0, *val;
             val = stats.lookup_or_init(&key, &value);
-            *val += 1;
+            if (val) {
+                *val += 1;
+            }
             return 0;
         }
         """
@@ -94,8 +98,10 @@ def test_struct_custom_func(self):
             u32 key=0;
             counter value = {0,0}, *val;
             val = stats.lookup_or_init(&key, &value);
-            val->c1 += 1;
-            val->c2 += 1;
+            if (val) {
+                val->c1 += 1;
+                val->c2 += 1;
+            }
             return 0;
         }
         """

tests/python/test_stat1.c

@@ -48,8 +48,10 @@ int on_packet(struct __sk_buff *skb) {
     }
     struct IPLeaf zleaf = {0};
     struct IPLeaf *leaf = stats.lookup_or_init(&key, &zleaf);
-    lock_xadd(&leaf->rx_pkts, rx);
-    lock_xadd(&leaf->tx_pkts, tx);
+    if (leaf) {
+      lock_xadd(&leaf->rx_pkts, rx);
+      lock_xadd(&leaf->tx_pkts, tx);
+    }
   }
 
 EOP:

tests/python/test_trace2.c

@@ -8,6 +8,9 @@ BPF_HASH(stats, struct Ptr, struct Counters, 1024);
 int count_sched(struct pt_regs *ctx) {
   struct Ptr key = {.ptr = PT_REGS_PARM1(ctx)};
   struct Counters zleaf = {0};
-  stats.lookup_or_init(&key, &zleaf)->stat1++;
+  struct Counters *val = stats.lookup_or_init(&key, &zleaf);
+  if (val) {
+    val->stat1++;
+  }
   return 0;
 }

tests/python/test_trace2.py

@@ -19,7 +19,10 @@
   struct Counters zleaf;
 
   memset(&zleaf, 0, sizeof(zleaf));
-  stats.lookup_or_init(&key, &zleaf)->stat1++;
+  struct Counters *val = stats.lookup_or_init(&key, &zleaf);
+  if (val) {
+    val->stat1++;
+  }
   return 0;
 }
 """

tests/python/test_trace3.c

@@ -48,6 +48,8 @@ int probe_blk_update_request(struct pt_regs *ctx) {
 
   u64 zero = 0;
   u64 *val = latency.lookup_or_init(&index, &zero);
-  lock_xadd(val, 1);
+  if (val) {
+    lock_xadd(val, 1);
+  }
   return 0;
 }

tests/python/test_trace4.py

@@ -14,11 +14,17 @@ def setUp(self):
         typedef struct { u64 val; } Val;
         BPF_HASH(stats, Key, Val, 3);
         int hello(void *ctx) {
-          stats.lookup_or_init(&(Key){1}, &(Val){0})->val++;
+          Val *val = stats.lookup_or_init(&(Key){1}, &(Val){0});
+          if (val) {
+            val->val++;
+          }
           return 0;
         }
         int goodbye(void *ctx) {
-          stats.lookup_or_init(&(Key){2}, &(Val){0})->val++;
+          Val *val = stats.lookup_or_init(&(Key){2}, &(Val){0});
+          if (val) {
+            val->val++;
+          }
           return 0;
         }
         """)

tests/python/test_trace_maxactive.py

@@ -14,11 +14,17 @@ def setUp(self):
         typedef struct { u64 val; } Val;
         BPF_HASH(stats, Key, Val, 3);
         int hello(void *ctx) {
-          stats.lookup_or_init(&(Key){1}, &(Val){0})->val++;
+          Val *val = stats.lookup_or_init(&(Key){1}, &(Val){0});
+          if (val) {
+            val->val++;
+          }
           return 0;
         }
         int goodbye(void *ctx) {
-          stats.lookup_or_init(&(Key){2}, &(Val){0})->val++;
+          Val *val = stats.lookup_or_init(&(Key){2}, &(Val){0});
+          if (val) {
+            val->val++;
+          }
           return 0;
         }
         """)