[CLIENT-5298] Address {Query,Scan}.results() and Scan.foreach() potentially missing records during a paginated query disruption

src/include/foreach.h

@@ -0,0 +1,11 @@
+#include <Python.h>
+#include "client.h"
+
+// Struct for Python User-Data for the Callback
+typedef struct {
+    PyObject *py_obj;
+    AerospikeClient *client;
+    as_vector thread_errors;
+    pthread_mutex_t thread_errors_mutex;
+    bool partition_query;
+} LocalData;

src/include/query.h

@@ -148,6 +148,11 @@ PyObject *StoreUnicodePyObject(AerospikeQuery *self, PyObject *obj);
 
 int64_t pyobject_to_int64(PyObject *py_obj);
 
+PyObject *AerospikeQuery_Foreach_Invoke(AerospikeQuery *self,
+                                        PyObject *py_callback,
+                                        PyObject *py_policy,
+                                        PyObject *py_options);
+
 // We need to share this with src/main/client/query.c because this function is no longer assigned
 // to the query type's tp_new slot. We are trying to prevent users from using the query type's constructor directly
 // to create a query instance.

src/include/scan.h

@@ -132,3 +132,9 @@ PyObject *AerospikeScan_Get_Partitions_status(AerospikeScan *self);
 
 AerospikeScan *AerospikeScan_Type_New(PyTypeObject *type,
                                       AerospikeClient *client);
+
+PyObject *AerospikeScan_Foreach_Invoke(AerospikeScan *self,
+                                       PyObject *py_callback,
+                                       PyObject *py_policy,
+                                       PyObject *py_options,
+                                       PyObject *py_nodename);

src/include/types.h

@@ -84,6 +84,13 @@ typedef struct {
 typedef struct {
     PyObject_HEAD AerospikeClient *client;
     as_query query;
+
+    // This is for query.results() to restore its original partitions status before being called
+    // We don't want to heap allocate a new partitions status every time we call query.results().
+    // And the stack size may be too small to hold the partitions status.
+    as_partitions_status *partitions_status_backup_buffer;
+    size_t partitions_status_backup_buffer_capacity;
+
     UnicodePyObjects u_objs;
     as_vector *unicodeStrVector;
     as_static_pool *static_pool;

src/main/convert_partition_filter.c

@@ -180,7 +180,7 @@ as_status convert_partition_filter(AerospikeClient *self,
         parts_setup(filter->begin, filter->count, //cluster->n_partitions,
                     &filter->digest);
 
-    if (parts_stat && PyDict_Check(parts_stat)) {
+    if (parts_stat && PyDict_Check(parts_stat) && PyDict_Size(parts_stat) > 0) {
 
         PyObject *py_done =
             PyDict_GetItemString(parts_stat, PARTITIONS_STATUS_KEY_DONE);

src/main/query/foreach.c

@@ -30,17 +30,9 @@
 #include "exceptions.h"
 #include "query.h"
 #include "policy.h"
+#include "foreach.h"
 
-// Struct for Python User-Data for the Callback
-typedef struct {
-    PyObject *callback;
-    AerospikeClient *client;
-    int partition_query;
-    as_vector thread_errors;
-    pthread_mutex_t thread_errors_mutex;
-} LocalData;
-
-static bool each_result(const as_val *val, void *udata)
+bool each_result(const as_val *val, void *udata)
 {
     bool retval = true;
 
@@ -50,7 +42,7 @@ static bool each_result(const as_val *val, void *udata)
 
     // Extract callback user-data
     LocalData *data = (LocalData *)udata;
-    PyObject *py_callback = data->callback;
+    PyObject *py_callback_or_list_of_results = data->py_obj;
 
     // Python Function Arguments and Result Value
     PyObject *py_arglist = NULL;
@@ -72,45 +64,61 @@ static bool each_result(const as_val *val, void *udata)
         goto EXIT_CALLBACK;
     }
 
-    // Build Python Function Arguments
-    if (data->partition_query) {
+    if (PyList_Check(py_callback_or_list_of_results)) {
+        // query.results()
+        if (py_result) {
+            int retval =
+                PyList_Append(py_callback_or_list_of_results, py_result);
+            Py_DECREF(py_result);
+            if (retval == -1) {
+                as_error_update(&thread_err_local, AEROSPIKE_ERR_CLIENT,
+                                "Failed to append item to results list");
+                goto EXIT_CALLBACK;
+            }
+        }
+    }
+    else {
+        // Build Python Function Arguments
+        if (data->partition_query) {
 
-        uint32_t part_id = 0;
+            uint32_t part_id = 0;
 
-        as_record *rec = as_record_fromval(val);
+            as_record *rec = as_record_fromval(val);
 
-        if (rec->key.digest.init) {
-            part_id =
-                as_partition_getid(rec->key.digest.value, CLUSTER_NPARTITIONS);
-        }
+            if (rec->key.digest.init) {
+                part_id = as_partition_getid(rec->key.digest.value,
+                                             CLUSTER_NPARTITIONS);
+            }
 
-        py_arglist = PyTuple_New(2);
+            py_arglist = PyTuple_New(2);
 
-        PyTuple_SetItem(py_arglist, 0, PyLong_FromUnsignedLong(part_id));
-        PyTuple_SetItem(py_arglist, 1, py_result);
-    }
-    else {
-        py_arglist = PyTuple_New(1);
-        PyTuple_SetItem(py_arglist, 0, py_result);
-    }
+            PyTuple_SetItem(py_arglist, 0, PyLong_FromUnsignedLong(part_id));
+            PyTuple_SetItem(py_arglist, 1, py_result);
+        }
+        else {
+            py_arglist = PyTuple_New(1);
+            PyTuple_SetItem(py_arglist, 0, py_result);
+        }
 
-    // Invoke Python Callback
-    py_return = PyObject_Call(py_callback, py_arglist, NULL);
-
-    // Release Python Function Arguments
-    Py_DECREF(py_arglist);
-    // handle return value
-    if (!py_return) {
-        // an exception was raised, handle it (someday)
-        // for now, we bail from the loop
-        as_error_update(&thread_err_local, AEROSPIKE_ERR_CLIENT,
-                        "Callback function contains an error");
-        retval = false;
-    }
-    else if (py_return == Py_False) {
-        retval = false;
+        // Invoke Python Callback
+        py_return =
+            PyObject_Call(py_callback_or_list_of_results, py_arglist, NULL);
+
+        // Release Python Function Arguments
+        Py_DECREF(py_arglist);
+        // handle return value
+        if (!py_return) {
+            // an exception was raised, handle it (someday)
+            // for now, we bail from the loop
+            as_error_update(&thread_err_local, AEROSPIKE_ERR_CLIENT,
+                            "Callback function contains an error");
+            retval = false;
+        }
+        else if (py_return == Py_False) {
+            retval = false;
+        }
+        Py_XDECREF(py_return);
     }
-    Py_XDECREF(py_return);
 
 EXIT_CALLBACK:
     if (thread_err_local.code != AEROSPIKE_OK) {
@@ -129,27 +137,13 @@ static bool each_result(const as_val *val, void *udata)
     return retval;
 }
 
-PyObject *AerospikeQuery_Foreach(AerospikeQuery *self, PyObject *args,
-                                 PyObject *kwds)
+PyObject *AerospikeQuery_Foreach_Invoke(AerospikeQuery *self,
+                                        PyObject *py_callback,
+                                        PyObject *py_policy,
+                                        PyObject *py_options)
 {
-    // Python Function Arguments
-    PyObject *py_callback = NULL;
-    PyObject *py_policy = NULL;
-    PyObject *py_options = NULL;
-    // Python Function Keyword Arguments
-    static char *kwlist[] = {"callback", "policy", "options", NULL};
-
-    // Python Function Argument Parsing
-    if (PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:foreach", kwlist,
-                                    &py_callback, &py_policy,
-                                    &py_options) == false) {
-        as_query_destroy(&self->query);
-        return NULL;
-    }
-
     // Initialize callback user data
-    LocalData data;
-    data.callback = py_callback;
+    LocalData data = {0};
     data.client = self->client;
     data.partition_query = 0;
 
@@ -171,7 +165,12 @@ PyObject *AerospikeQuery_Foreach(AerospikeQuery *self, PyObject *args,
     as_partition_filter *partition_filter_p = NULL;
     as_partitions_status *ps = NULL;
 
-    // Initialize error
+    bool is_query_results = py_callback == NULL;
+
+    as_partitions_status *backed_up_part_status = NULL;
+    // query.results() may raise an exception before the backup is made
+    // so we don't want to reset the query's current partitions status in this case
+    bool is_query_state_backed_up = false;
 
     if (!self || !self->client->as) {
         as_error_update(&err, AEROSPIKE_ERR_PARAM, "Invalid aerospike object");
@@ -184,6 +183,18 @@ PyObject *AerospikeQuery_Foreach(AerospikeQuery *self, PyObject *args,
         goto CLEANUP;
     }
 
+    if (is_query_results) {
+        data.py_obj = PyList_New(0);
+        if (data.py_obj == NULL) {
+            as_error_update(&err, AEROSPIKE_ERR_CLIENT,
+                            "Was unable to construct results list");
+            goto CLEANUP;
+        }
+    }
+    else {
+        data.py_obj = py_callback;
+    }
+
     // Convert python policy object to as_policy_exists
     pyobject_to_policy_query(
         self->client, &err, py_policy, &query_policy, &query_policy_p,
@@ -212,6 +223,50 @@ PyObject *AerospikeQuery_Foreach(AerospikeQuery *self, PyObject *args,
         goto CLEANUP;
     }
 
+    if (is_query_results) {
+        as_partitions_status *backup_source = NULL;
+        if (self->query.parts_all) {
+            // This query is resuming.
+            // If there is a user-provided partitions status, the C client will still ignore it in this case.
+            backup_source = self->query.parts_all;
+        }
+        else if (ps) {
+            // This is a new query object
+            backup_source = ps;
+        }
+
+        if (backup_source) {
+            size_t backup_size_bytes =
+                sizeof(as_partitions_status) +
+                sizeof(as_partition_status) * backup_source->part_count;
+
+            // The number of partitions is always 4096, but we have this code in case
+            // self->partitions_status_backup_buffer was never initialized before
+            if (self->partitions_status_backup_buffer_capacity <
+                backup_size_bytes) {
+                as_partitions_status *new_buffer = cf_malloc(backup_size_bytes);
+                if (!new_buffer) {
+                    as_error_update(&err, AEROSPIKE_ERR_CLIENT,
+                                    "Failed to back up partitions status");
+                    goto CLEANUP;
+                }
+
+                if (self->partitions_status_backup_buffer) {
+                    as_partitions_status_release(
+                        self->partitions_status_backup_buffer);
+                }
+                self->partitions_status_backup_buffer = new_buffer;
+                self->partitions_status_backup_buffer_capacity =
+                    backup_size_bytes;
+            }
+
+            backed_up_part_status = self->partitions_status_backup_buffer;
+            memcpy(backed_up_part_status, backup_source, backup_size_bytes);
+            backed_up_part_status->ref_count = 1;
+        }
+        is_query_state_backed_up = true;
+    }
+
     Py_BEGIN_ALLOW_THREADS
 
     // Invoke operation
@@ -260,10 +315,53 @@ PyObject *AerospikeQuery_Foreach(AerospikeQuery *self, PyObject *args,
     pthread_mutex_destroy(&data.thread_errors_mutex);
 
     if (err.code != AEROSPIKE_OK) {
-        raise_exception_base(&err, Py_None, Py_None, Py_None, Py_None, Py_None);
+        if (is_query_results) {
+            if (is_query_state_backed_up) {
+                if (self->query.parts_all) {
+                    as_partitions_status_release(self->query.parts_all);
+                }
+                self->query.parts_all = backed_up_part_status;
+                if (backed_up_part_status) {
+                    // query.parts_all is now the owner of the backup, not the query's backup buffer ptr
+                    // If the query object is destroyed, we don't want to double free the partitions status
+                    self->partitions_status_backup_buffer = NULL;
+                    self->partitions_status_backup_buffer_capacity = 0;
+                }
+            }
+
+            Py_XDECREF(data.py_obj);
+        }
+
+        raise_exception(&err);
+        return NULL;
+    }
+
+    if (is_query_results) {
+        return data.py_obj;
+    }
+    else {
+        Py_RETURN_NONE;
+    }
+}
+
+PyObject *AerospikeQuery_Foreach(AerospikeQuery *self, PyObject *args,
+                                 PyObject *kwds)
+{
+    // Python Function Arguments
+    PyObject *py_callback = NULL;
+    PyObject *py_policy = NULL;
+    PyObject *py_options = NULL;
+    // Python Function Keyword Arguments
+    static char *kwlist[] = {"callback", "policy", "options", NULL};
+
+    // Python Function Argument Parsing
+    if (PyArg_ParseTupleAndKeywords(args, kwds, "O|OO:foreach", kwlist,
+                                    &py_callback, &py_policy,
+                                    &py_options) == false) {
+        as_query_destroy(&self->query);
         return NULL;
     }
 
-    Py_INCREF(Py_None);
-    return Py_None;
+    return AerospikeQuery_Foreach_Invoke(self, py_callback, py_policy,
+                                         py_options);
 }