[PATCH v2] Add new built-in tapset function quit()

["Yichun Zhang (agentzh)" <agentzh at gmail dot com>]

The new built-in tapset function quit() is similar to exit(), but it
aborts the currently executing functions and probe handlers immediately.
Works in both the kernel and dyninst runtimes. The bpf runtime is not
yet supported.

fche thinks it is already too late to change the current behavior of
exit(), hence this new function.

fche suggests the global variable name _stp_abort_flag is good since
it looks alarming and is consistent with the existing global variable
_stp_exit_flag.

Also added corresponding tests for both quit() and exit().

This new function can be disabled by the '--compatible 3.3' option. Also
added tests for this.

Updated NEWS for this new feature.
---
 NEWS                                |   6 +
 runtime/dyninst/io.c                |   9 ++
 runtime/linux/io.c                  |  13 +++
 tapset/logging.stp                  |  23 ++++
 testsuite/systemtap.base/exit.exp   | 112 +++++++++++++++++++
 testsuite/systemtap.base/exit_1.stp |  11 ++
 testsuite/systemtap.base/exit_2.stp |  17 +++
 testsuite/systemtap.base/exit_3.stp |   5 +
 testsuite/systemtap.base/quit.exp   | 214 ++++++++++++++++++++++++++++++++++++
 testsuite/systemtap.base/quit_1.stp |  11 ++
 testsuite/systemtap.base/quit_2.stp |  17 +++
 testsuite/systemtap.base/quit_3.stp |   5 +
 translate.cxx                       |   4 +-
 13 files changed, 445 insertions(+), 2 deletions(-)
 create mode 100644 testsuite/systemtap.base/exit.exp
 create mode 100644 testsuite/systemtap.base/exit_1.stp
 create mode 100644 testsuite/systemtap.base/exit_2.stp
 create mode 100644 testsuite/systemtap.base/exit_3.stp
 create mode 100644 testsuite/systemtap.base/quit.exp
 create mode 100644 testsuite/systemtap.base/quit_1.stp
 create mode 100644 testsuite/systemtap.base/quit_2.stp
 create mode 100644 testsuite/systemtap.base/quit_3.stp

diff --git a/NEWS b/NEWS
index 7ed02a99f..f1ba50921 100644
--- a/NEWS
+++ b/NEWS
@@ -64,6 +64,12 @@
   (The syscall tapsets are broken on kernel 4.17-rc, and will be fixed
   in a next release coming soon; PR23160.)
 
+- Add new built-in tapset function quit() which is similar to exit(),
+  but it aborts the currently executing functions and probe handlers
+  immediately instead of waiting. Works in both the kernel and dyninst
+  runtimes. This function can be disabled by the '--compatible 3.3'
+  option.
+
 * What's new in version 3.2, 2017-10-18
 
 - SystemTap now includes an extended Berkeley Packet Filter (eBPF)
diff --git a/runtime/dyninst/io.c b/runtime/dyninst/io.c
index b8e9568a5..6070597ac 100644
--- a/runtime/dyninst/io.c
+++ b/runtime/dyninst/io.c
@@ -19,6 +19,9 @@
 #define WARN_STRING "WARNING: "
 #define ERR_STRING "ERROR: "
 
+// a global flag to signal immediate probe/function execution abort from quit().
+static int _stp_abort_flag = 0;
+
 enum code { INFO=0, WARN, ERROR, DBUG };
 
 static void _stp_vlog (enum code type, const char *func, int line,
@@ -165,4 +168,10 @@ static void _stp_exit (void)
 	_stp_dyninst_transport_request_exit();
 }
 
+static void _stp_quit (void)
+{
+	_stp_abort_flag = 1;
+	_stp_exit ();
+}
+
 #endif /* _STAPDYN_IO_C_ */
diff --git a/runtime/linux/io.c b/runtime/linux/io.c
index 74a032c52..22245a666 100644
--- a/runtime/linux/io.c
+++ b/runtime/linux/io.c
@@ -24,6 +24,9 @@
 #error "STP_LOG_BUF_LEN is too short"
 #endif
 
+// a global flag to signal immediate probe/function execution abort from quit().
+static int _stp_abort_flag = 0;
+
 enum code { INFO=0, WARN, ERROR, DBUG };
 
 static void _stp_vlog (enum code type, const char *func, int line, const char *fmt, va_list args)
@@ -105,6 +108,16 @@ static void _stp_exit (void)
 	_stp_exit_flag = 1;
 }
 
+/** Similar to _stp_exit(), but also sets the flag
+ * _stp_abort_flag.
+ */
+static void _stp_quit (void)
+{
+	_stp_abort_flag = 1;
+	_stp_exit ();
+}
+
+
 /** Prints error message and exits.
  * This function sends an error message immediately to staprun. It
  * will also be sent over the bulk transport (relayfs) if it is
diff --git a/tapset/logging.stp b/tapset/logging.stp
index d09b8eb80..352f0fbe8 100644
--- a/tapset/logging.stp
+++ b/tapset/logging.stp
@@ -74,6 +74,29 @@ function exit ()
 %)
 
 
+%(systemtap_v >= "4.0" %?
+/**
+ * sfunction quit - Immediately shutting down probing script.
+ *
+ * Description: This is similar to exit() but immediately stops
+ * the currently executing functions and probes.
+ */
+function quit ()
+%( runtime != "bpf" %?
+  %{ /* unprivileged */
+     atomic_set (session_state(), STAP_SESSION_STOPPING);
+     _stp_quit ();
+  %}
+%:
+  { /* unprivileged */ /* bpf */
+    _set_exit_status()
+    printf("ERROR: quit() not supported yet\n")
+    exit()  /* TODO: need to abort the execution flow immediately */
+  }
+%)
+%)
+
+
 /**
  * sfunction error - Send an error message
  *
diff --git a/testsuite/systemtap.base/exit.exp b/testsuite/systemtap.base/exit.exp
new file mode 100644
index 000000000..9e91c30ec
--- /dev/null
+++ b/testsuite/systemtap.base/exit.exp
@@ -0,0 +1,112 @@
+set test "exit"
+set testpath "$srcdir/$subdir"
+
+if {! [installtest_p]} { untested "$test"; return }
+
+# --- TEST 1 ---
+
+set subtest1 "TEST 1: exit() in the middle of a func body"
+foreach runtime [get_runtime_list] {
+    if {$runtime eq ""} {
+        set runtime "kernel"
+    }
+    set test_name "$test: $subtest1 ($runtime)"
+
+    set cmd "stap --runtime=$runtime '$srcdir/$subdir/${test}_1.stp'"
+    send_log "executing: $cmd\n"
+    set pipe [open "| sh -c {$cmd}" r]
+    set out [read $pipe]
+    set failed 0
+    if {[catch {close $pipe} stderr] != 0} {
+        set failed 1
+    }
+
+    set exp_out "enter f\nleave f\nexit\n"
+    regsub -all -- {\n} $exp_out {\n} escaped_exp_out
+    if {$out eq $exp_out} {
+        pass "${test_name}: stdout matches \"$escaped_exp_out\""
+    } else {
+        fail "${test_name}: stdout fails to match \"$escaped_exp_out\": got \"$out\""
+    }
+
+    if {$failed} {
+        fail "${test_name}: exit code not zero"
+    } else {
+        pass "${test_name}: exit code is zero"
+    }
+    if {$stderr ne ""} {
+        send_log "stderr:\n$stderr"
+    }
+}
+
+# --- TEST 2 ---
+
+set subtest2 "TEST 2: exit() in the middle of a func body in a deeper func call chain"
+foreach runtime [get_runtime_list] {
+    if {$runtime eq ""} {
+        set runtime "kernel"
+    }
+    set test_name "$test: $subtest2 ($runtime)"
+
+    set cmd "stap --runtime=$runtime '$srcdir/$subdir/${test}_2.stp'"
+    send_log "executing: $cmd\n"
+    set pipe [open "| sh -c {$cmd}" r]
+    set out [read $pipe]
+    set failed 0
+    if {[catch {close $pipe} stderr] != 0} {
+        set failed 1
+    }
+
+    set exp_out "enter g\nenter f\nleave f\nleave g\nexit\n"
+    regsub -all -- {\n} $exp_out {\n} escaped_exp_out
+    if {$out eq $exp_out} {
+        pass "${test_name}: stdout matches \"$escaped_exp_out\""
+    } else {
+        fail "${test_name}: stdout fails to match \"$escaped_exp_out\": got \"$out\""
+    }
+
+    if {$failed} {
+        fail "${test_name}: exit code not zero"
+    } else {
+        pass "${test_name}: exit code is zero"
+    }
+    if {$stderr ne ""} {
+        send_log "stderr:\n$stderr"
+    }
+}
+
+# --- TEST 3 ---
+
+set subtest3 "TEST 3: exit() in the middle of a probe handler body"
+foreach runtime [get_runtime_list] {
+    if {$runtime eq ""} {
+        set runtime "kernel"
+    }
+    set test_name "$test: $subtest3 ($runtime)"
+
+    set cmd "stap --runtime=$runtime '$srcdir/$subdir/${test}_3.stp'"
+    send_log "executing: $cmd\n"
+    set pipe [open "| sh -c {$cmd}" r]
+    set out [read $pipe]
+    set failed 0
+    if {[catch {close $pipe} stderr] != 0} {
+        set failed 1
+    }
+
+    set exp_out "enter probe\nleave probe\n"
+    regsub -all -- {\n} $exp_out {\n} escaped_exp_out
+    if {$out eq $exp_out} {
+        pass "${test_name}: stdout matches \"$escaped_exp_out\""
+    } else {
+        fail "${test_name}: stdout fails to match \"$escaped_exp_out\": got \"$out\""
+    }
+
+    if {$failed} {
+        fail "${test_name}: exit code not zero"
+    } else {
+        pass "${test_name}: exit code is zero"
+    }
+    if {$stderr ne ""} {
+        send_log "stderr:\n$stderr"
+    }
+}
diff --git a/testsuite/systemtap.base/exit_1.stp b/testsuite/systemtap.base/exit_1.stp
new file mode 100644
index 000000000..c2dd0141a
--- /dev/null
+++ b/testsuite/systemtap.base/exit_1.stp
@@ -0,0 +1,11 @@
+function f() {
+    println("enter f");
+    exit();
+    println("leave f");
+}
+
+probe begin {
+    f();
+    printf("exit\n");
+    exit();
+}
diff --git a/testsuite/systemtap.base/exit_2.stp b/testsuite/systemtap.base/exit_2.stp
new file mode 100644
index 000000000..b747ddc0c
--- /dev/null
+++ b/testsuite/systemtap.base/exit_2.stp
@@ -0,0 +1,17 @@
+function f() {
+    println("enter f");
+    exit();
+    println("leave f");
+}
+
+function g() {
+    println("enter g");
+    f();
+    println("leave g");
+}
+
+probe begin {
+    g();
+    printf("exit\n");
+    exit();
+}
diff --git a/testsuite/systemtap.base/exit_3.stp b/testsuite/systemtap.base/exit_3.stp
new file mode 100644
index 000000000..3f3b801d1
--- /dev/null
+++ b/testsuite/systemtap.base/exit_3.stp
@@ -0,0 +1,5 @@
+probe begin {
+    println("enter probe");
+    exit();
+    println("leave probe");
+}
diff --git a/testsuite/systemtap.base/quit.exp b/testsuite/systemtap.base/quit.exp
new file mode 100644
index 000000000..e02be5ad0
--- /dev/null
+++ b/testsuite/systemtap.base/quit.exp
@@ -0,0 +1,214 @@
+set test "quit"
+set testpath "$srcdir/$subdir"
+
+if {! [installtest_p]} { untested "$test"; return }
+
+# --- TEST 1 ---
+
+set subtest1 "TEST 1: quit() in the middle of a func body"
+foreach runtime [get_runtime_list] {
+    if {$runtime eq ""} {
+        set runtime "kernel"
+    }
+    set test_name "$test: $subtest1 ($runtime)"
+
+    set cmd "stap --runtime=$runtime '$srcdir/$subdir/${test}_1.stp'"
+    send_log "executing: $cmd\n"
+    set pipe [open "| sh -c {$cmd}" r]
+    set out [read $pipe]
+    set failed 0
+    if {[catch {close $pipe} stderr] != 0} {
+        if {$stderr ne "" && [string index $stderr end] ne "\n"} {
+            append stderr "\n"
+        }
+        global errorCode
+        if {"CHILDSTATUS" == [lindex $errorCode 0]} {
+            set failed [lindex $errorCode 2]
+        }
+    }
+
+    set exp_out "enter f\n"
+    regsub -all -- {\n} $exp_out {\n} escaped_exp_out
+    if {$out eq $exp_out} {
+        pass "${test_name}: stdout matches \"$escaped_exp_out\""
+    } else {
+        fail "${test_name}: stdout fails to match \"$escaped_exp_out\": got \"$out\""
+    }
+
+    if {$failed} {
+        fail "${test_name}: exit code should be zero but is $failed"
+    } else {
+        pass "${test_name}: exit code is zero"
+    }
+    if {$stderr ne ""} {
+        send_log "stderr:\n$stderr"
+    }
+}
+
+# --- TEST 2 ---
+
+set subtest2 "TEST 2: quit() in the middle of a func body in a deeper func call chain"
+foreach runtime [get_runtime_list] {
+    if {$runtime eq ""} {
+        set runtime "kernel"
+    }
+    set test_name "$test: $subtest2 ($runtime)"
+
+    set cmd "stap --runtime=$runtime '$srcdir/$subdir/${test}_2.stp'"
+    send_log "executing: $cmd\n"
+    set pipe [open "| sh -c {$cmd}" r]
+    set out [read $pipe]
+    set failed 0
+    if {[catch {close $pipe} stderr] != 0} {
+        if {$stderr ne "" && [string index $stderr end] ne "\n"} {
+            append stderr "\n"
+        }
+        global errorCode
+        if {"CHILDSTATUS" == [lindex $errorCode 0]} {
+            set failed [lindex $errorCode 2]
+        }
+    }
+
+    set exp_out "enter g\nenter f\n"
+    regsub -all -- {\n} $exp_out {\n} escaped_exp_out
+    if {$out eq $exp_out} {
+        pass "${test_name}: stdout matches \"$escaped_exp_out\""
+    } else {
+        fail "${test_name}: stdout fails to match \"$escaped_exp_out\": got \"$out\""
+    }
+
+    if {$failed} {
+        fail "${test_name}: exit code should be zero but is $failed"
+    } else {
+        pass "${test_name}: exit code is zero"
+    }
+    if {$stderr ne ""} {
+        send_log "stderr:\n$stderr"
+    }
+}
+
+# --- TEST 3 ---
+
+set subtest3 "TEST 3: quit() in the middle of a probe handler body"
+foreach runtime [get_runtime_list] {
+    if {$runtime eq ""} {
+        set runtime "kernel"
+    }
+    set test_name "$test: $subtest3 ($runtime)"
+
+    set cmd "stap --runtime=$runtime '$srcdir/$subdir/${test}_3.stp'"
+    send_log "executing: $cmd\n"
+    set pipe [open "| sh -c {$cmd}" r]
+    set out [read $pipe]
+    set failed 0
+    if {[catch {close $pipe} stderr] != 0} {
+        if {$stderr ne "" && [string index $stderr end] ne "\n"} {
+            append stderr "\n"
+        }
+        global errorCode
+        if {"CHILDSTATUS" == [lindex $errorCode 0]} {
+            set failed [lindex $errorCode 2]
+        }
+    }
+
+    set exp_out "enter probe\n"
+    regsub -all -- {\n} $exp_out {\n} escaped_exp_out
+    if {$out eq $exp_out} {
+        pass "${test_name}: stdout matches \"$escaped_exp_out\""
+    } else {
+        fail "${test_name}: stdout fails to match \"$escaped_exp_out\": got \"$out\""
+    }
+
+    if {$failed} {
+        fail "${test_name}: exit code should be zero but is $failed"
+    } else {
+        pass "${test_name}: exit code is zero"
+    }
+    if {$stderr ne ""} {
+        send_log "stderr:\n$stderr"
+    }
+}
+
+# --- TEST 4 ---
+
+set subtest4 "TEST 4: quit() in the middle of a probe handler body (--compatible 4.0)"
+foreach runtime [get_runtime_list] {
+    if {$runtime eq ""} {
+        set runtime "kernel"
+    }
+    set test_name "$test: $subtest4 ($runtime)"
+
+    set cmd "stap --compatible 4.0 --runtime=$runtime '$srcdir/$subdir/${test}_3.stp'"
+    send_log "executing: $cmd\n"
+    set pipe [open "| sh -c {$cmd}" r]
+    set out [read $pipe]
+    set failed 0
+    if {[catch {close $pipe} stderr] != 0} {
+        if {$stderr ne "" && [string index $stderr end] ne "\n"} {
+            append stderr "\n"
+        }
+        global errorCode
+        if {"CHILDSTATUS" == [lindex $errorCode 0]} {
+            set failed [lindex $errorCode 2]
+        }
+    }
+
+    set exp_out "enter probe\n"
+    regsub -all -- {\n} $exp_out {\n} escaped_exp_out
+    if {$out eq $exp_out} {
+        pass "${test_name}: stdout matches \"$escaped_exp_out\""
+    } else {
+        fail "${test_name}: stdout fails to match \"$escaped_exp_out\": got \"$out\""
+    }
+
+    if {$failed} {
+        fail "${test_name}: exit code should be zero but is $failed"
+    } else {
+        pass "${test_name}: exit code is zero"
+    }
+    if {$stderr ne ""} {
+        send_log "stderr:\n$stderr"
+    }
+}
+
+# --- TEST 5 ---
+
+set subtest5 "TEST 5: quit() in the middle of a probe handler body (--compatible 3.3)"
+set test_name "$test: $subtest5"
+
+set cmd "stap --compatible 3.3 '$srcdir/$subdir/${test}_3.stp'"
+send_log "executing: $cmd\n"
+set pipe [open "| sh -c {$cmd}" r]
+set out [read $pipe]
+set failed 0
+if {[catch {close $pipe} stderr] != 0} {
+    if {$stderr ne "" && [string index $stderr end] ne "\n"} {
+        append stderr "\n"
+    }
+    global errorCode
+    if {"CHILDSTATUS" == [lindex $errorCode 0]} {
+        set failed [lindex $errorCode 2]
+    }
+}
+
+set exp_out ""
+regsub -all -- {\n} $exp_out {\n} escaped_exp_out
+if {$out eq $exp_out} {
+    pass "${test_name}: stdout matches \"$escaped_exp_out\""
+} else {
+    fail "${test_name}: stdout fails to match \"$escaped_exp_out\": got \"$out\""
+}
+
+if {$failed} {
+    pass "${test_name}: exit code should be non-zero"
+} else {
+    fail "${test_name}: exit code should be non-zero but is zero"
+}
+
+set stderr_pat "^semantic error: unresolved function \\(similar: \[^\\n\]*?\\): identifier 'quit' at \[^\\n\]*?\\.stp:3:5\\n"
+regsub -all -- {\n} $stderr_pat {\n} escaped_stderr_pat
+if {[regexp -linestop -lineanchor -- $stderr_pat $stderr]} {
+    pass "${test_name}: stderr matches \"$escaped_stderr_pat\""
+} else {
+    fail "${test_name}: stderr fails to match \"$escaped_stderr_pat\": got \"$stderr\""
+}
diff --git a/testsuite/systemtap.base/quit_1.stp b/testsuite/systemtap.base/quit_1.stp
new file mode 100644
index 000000000..f5903db45
--- /dev/null
+++ b/testsuite/systemtap.base/quit_1.stp
@@ -0,0 +1,11 @@
+function f() {
+    println("enter f");
+    quit();
+    println("leave f");
+}
+
+probe oneshot {
+    f();
+    printf("quit\n");
+    quit();
+}
diff --git a/testsuite/systemtap.base/quit_2.stp b/testsuite/systemtap.base/quit_2.stp
new file mode 100644
index 000000000..a4d0d7caf
--- /dev/null
+++ b/testsuite/systemtap.base/quit_2.stp
@@ -0,0 +1,17 @@
+function f() {
+    println("enter f");
+    quit();
+    println("leave f");
+}
+
+function g() {
+    println("enter g");
+    f();
+    println("leave g");
+}
+
+probe oneshot {
+    g();
+    printf("quit\n");
+    quit();
+}
diff --git a/testsuite/systemtap.base/quit_3.stp b/testsuite/systemtap.base/quit_3.stp
new file mode 100644
index 000000000..6782404ce
--- /dev/null
+++ b/testsuite/systemtap.base/quit_3.stp
@@ -0,0 +1,5 @@
+probe oneshot {
+    println("enter probe");
+    quit();
+    println("leave probe");
+}
diff --git a/translate.cxx b/translate.cxx
index c7e7aa53f..743d6cd0e 100644
--- a/translate.cxx
+++ b/translate.cxx
@@ -5732,7 +5732,7 @@ c_unparser::visit_functioncall (functioncall* e)
 
       // call function
       o->newline() << c_funcname (r->name) << " (c);";
-      o->newline() << "if (unlikely(c->last_error)) goto out;";
+      o->newline() << "if (unlikely(c->last_error || _stp_abort_flag)) goto out;";
 
       if (!already_checked_action_count && !session->suppress_time_limits
           && !session->unoptimized)
@@ -6069,7 +6069,7 @@ c_unparser::visit_print_format (print_format* e)
 	}
       o->line() << ");";
       o->newline(-1) << "#endif // STP_LEGACY_PRINT";
-      o->newline() << "if (unlikely(c->last_error)) goto out;";
+      o->newline() << "if (unlikely(c->last_error || _stp_abort_flag)) goto out;";
       o->newline() << res.value() << ";";
     }
 }
-- 
2.11.0.295.gd7dffce