5 files changed, 15 insertions, 69 deletions
diff --git a/runtime/probes/bench/bench.c b/runtime/probes/bench/bench.c
index 9192d3f2..adbe01d4 100644
--- a/runtime/probes/bench/bench.c
+++ b/runtime/probes/bench/bench.c
@@ -35,16 +35,13 @@ static struct kprobe kp[] = {
 #define NUM_JPROBES (sizeof(jp)/sizeof(struct jprobe))
 #define NUM_KPROBES (sizeof(kp)/sizeof(struct kprobe))
 
-int init_module(void)
+int probe_start(void)
 {
   int ret;
-
-  TRANSPORT_OPEN;
-  
   ret = _stp_register_jprobes (jp, NUM_JPROBES);
   if (ret >= 0)
-    ret = _stp_register_kprobes (kp, NUM_KPROBES);
-  
+    if ((ret = _stp_register_kprobes (kp, NUM_KPROBES)) < 0)
+      _stp_unregister_jprobes (jp, NUM_JPROBES)      ;
   return ret;
 }
 
@@ -53,10 +50,3 @@ static void probe_exit (void)
   _stp_unregister_jprobes (jp, NUM_JPROBES); 
   _stp_unregister_kprobes (kp, NUM_KPROBES); 
 }
-
-void cleanup_module(void)
-{
-  _stp_transport_close();
-}
-
-MODULE_LICENSE("GPL");
diff --git a/runtime/probes/bench/bench_io1.c b/runtime/probes/bench/bench_io1.c
index 4cc890df..a0a44d34 100644
--- a/runtime/probes/bench/bench_io1.c
+++ b/runtime/probes/bench/bench_io1.c
@@ -36,14 +36,9 @@ static struct kprobe kp[] = {
 
 #define NUM_KPROBES (sizeof(kp)/sizeof(struct kprobe))
 
-int init_module(void)
+int probe_start(void)
 {
-  int ret;
-
-  TRANSPORT_OPEN;
-  
-  ret = _stp_register_kprobes (kp, NUM_KPROBES);
-  return ret;
+  return _stp_register_kprobes (kp, NUM_KPROBES);
 }
 
 static void probe_exit (void)
@@ -52,10 +47,3 @@ static void probe_exit (void)
   _stp_printf("dropped %d packets\n", atomic_read(&_stp_transport_failures));
   _stp_print_flush();
 }
-
-void cleanup_module(void)
-{
-  _stp_transport_close();
-}
-
-MODULE_LICENSE("GPL");
diff --git a/runtime/probes/bench/bench_io2.c b/runtime/probes/bench/bench_io2.c
index d606fa64..14ef747b 100644
--- a/runtime/probes/bench/bench_io2.c
+++ b/runtime/probes/bench/bench_io2.c
@@ -36,14 +36,9 @@ static struct kprobe kp[] = {
 
 #define NUM_KPROBES (sizeof(kp)/sizeof(struct kprobe))
 
-int init_module(void)
+int probe_start(void)
 {
-  int ret;
-
-  TRANSPORT_OPEN;
-  
-  ret = _stp_register_kprobes (kp, NUM_KPROBES);
-  return ret;
+  return _stp_register_kprobes (kp, NUM_KPROBES);
 }
 
 static void probe_exit (void)
@@ -51,10 +46,3 @@ static void probe_exit (void)
   _stp_unregister_kprobes (kp, NUM_KPROBES); 
   _stp_print_flush();
 }
-
-void cleanup_module(void)
-{
-  _stp_transport_close();
-}
-
-MODULE_LICENSE("GPL");
diff --git a/runtime/probes/bench/bench_multi.c b/runtime/probes/bench/bench_multi.c
index 566e8520..eab75878 100644
--- a/runtime/probes/bench/bench_multi.c
+++ b/runtime/probes/bench/bench_multi.c
@@ -62,24 +62,12 @@ static struct kprobe kp[] = {
 
 #define NUM_KPROBES (sizeof(kp)/sizeof(struct kprobe))
 
-int init_module(void)
+int probe_start(void)
 {
-  int ret;
-
-  TRANSPORT_OPEN;
-  
-  ret = _stp_register_kprobes (kp, NUM_KPROBES);
-  return ret;
+  return _stp_register_kprobes (kp, NUM_KPROBES);
 }
 
 static void probe_exit (void)
 {
   _stp_unregister_kprobes (kp, NUM_KPROBES); 
 }
-
-void cleanup_module(void)
-{
-  _stp_transport_close();
-}
-
-MODULE_LICENSE("GPL");
diff --git a/runtime/probes/bench/bench_ret.c b/runtime/probes/bench/bench_ret.c
index f759691d..d70f0ee1 100644
--- a/runtime/probes/bench/bench_ret.c
+++ b/runtime/probes/bench/bench_ret.c
@@ -44,14 +44,13 @@ static struct kprobe kp[] = {
 
 #define NUM_KPROBES (sizeof(kpr)/sizeof(struct kretprobe))
 
-int init_module(void)
+int probe_start(void)
 {
-  int ret;
-
-  TRANSPORT_OPEN;
-  
-  ret = _stp_register_kretprobes (kpr, NUM_KPROBES);
-  ret = _stp_register_kprobes (kp, 1);
+  int ret = _stp_register_kretprobes (kpr, NUM_KPROBES);
+  if (ret >= 0) {
+    if ((ret = _stp_register_kprobes (kp, 1)) < 0)
+      _stp_unregister_kretprobes (kpr, NUM_KPROBES);
+  }
   return ret;
 }
 
@@ -60,10 +59,3 @@ static void probe_exit (void)
   _stp_unregister_kretprobes (kpr, NUM_KPROBES); 
   _stp_unregister_kprobes (kp, 1); 
 }
-
-void cleanup_module(void)
-{
-  _stp_transport_close();
-}
-
-MODULE_LICENSE("GPL");