[PATCH 2/2] trace-cruncher: More options for stopping data iteration

["Yordan Karadzhov (VMware)" <y.karadz@gmail.com>]

So far the only option to stop the iteration over the tracing data
from the callback function was to call 'sys.exit()'. However, this
only stops the iteration without actually exiting. With this change
calling 'sys.exit()' will terminate (exit the module), while in
order to just stop the iteration, the callback should return nonzero
integer value.

Signed-off-by: Yordan Karadzhov (VMware) <y.karadz@gmail.com>
---
 src/ftracepy-utils.c | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/src/ftracepy-utils.c b/src/ftracepy-utils.c
index 7b1ff1b..814e1d6 100644
--- a/src/ftracepy-utils.c
+++ b/src/ftracepy-utils.c
@@ -2101,21 +2101,36 @@ static int callback(struct tep_event *event, struct tep_record *record,
 	ret = PyObject_CallObject((PyObject *)ctx->py_callback, arglist);
 	Py_DECREF(arglist);
 
-	if (ret) {
-		Py_DECREF(ret);
-	} else {
+	if (!ret) {
 		if (PyErr_Occurred()) {
-			if (PyErr_ExceptionMatches(PyExc_SystemExit)) {
-				PyErr_Clear();
+			PyObject *err_type, *err_value, *err_traceback;
+
+			PyErr_Fetch(&err_type, &err_value, &err_traceback);
+			if (err_type == PyExc_SystemExit) {
+				if (PyLong_CheckExact(err_value))
+					Py_Exit(PyLong_AsLong(err_value));
+				else
+					Py_Exit(0);
 			} else {
+				PyErr_Restore(err_type, err_value, err_traceback);
 				PyErr_Print();
 			}
 		}
 
-		ctx->status = false;
+		goto stop;
+	}
+
+	if (PyLong_CheckExact(ret) && PyLong_AsLong(ret) != 0) {
+		Py_DECREF(ret);
+		goto stop;
 	}
 
+	Py_DECREF(ret);
 	return 0;
+
+ stop:
+	ctx->status = false;
+	return 1;
 }
 
 static bool notrace_this_pid(struct tracefs_instance *instance)
-- 
2.30.2