don't segfault on low memory

Fixes #8

parquet/parquet.cc

@@ -72,6 +72,7 @@ static int parquetConnect(
   sqlite3_vtab **ppVtab,
   char **pzErr
 ){
+  try {
     if(argc != 4 || strlen(argv[3]) < 2) {
       *pzErr = sqlite3_mprintf("must provide exactly one argument, the path to a parquet file");
       return SQLITE_ERROR;
@@ -100,6 +101,11 @@ static int parquetConnect(
       *pzErr = sqlite3_mprintf(e.what());
       return SQLITE_ERROR;
     }
+  } catch(std::bad_alloc& ba) {
+    return SQLITE_NOMEM;
+  } catch(std::exception& e) {
+    return SQLITE_ERROR;
+  }
 }
 
 /*
@@ -131,6 +137,7 @@ static int parquetClose(sqlite3_vtab_cursor *cur){
 ** Constructor for a new sqlite3_vtab_parquet cursor object.
 */
 static int parquetOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor){
+  try {
     std::unique_ptr<sqlite3_vtab_cursor_parquet, void(*)(void*)> cursor(
         (sqlite3_vtab_cursor_parquet*)sqlite3_malloc(sizeof(sqlite3_vtab_cursor_parquet)),
         sqlite3_free);
@@ -141,6 +148,11 @@ static int parquetOpen(sqlite3_vtab *p, sqlite3_vtab_cursor **ppCursor){
 
     *ppCursor = (sqlite3_vtab_cursor*)cursor.release();
     return SQLITE_OK;
+  } catch(std::bad_alloc& ba) {
+    return SQLITE_NOMEM;
+  } catch(std::exception& e) {
+    return SQLITE_ERROR;
+  }
 }
 
 const char* opName(int op) {
@@ -183,9 +195,15 @@ const char* opName(int op) {
 ** Set the EOF marker if we reach the end of input.
 */
 static int parquetNext(sqlite3_vtab_cursor *cur){
+  try {
     ParquetCursor* cursor = ((sqlite3_vtab_cursor_parquet*)cur)->cursor;
     cursor->next();
     return SQLITE_OK;
+  } catch(std::bad_alloc& ba) {
+    return SQLITE_NOMEM;
+  } catch(std::exception& e) {
+    return SQLITE_ERROR;
+  }
 }
 
 /*
@@ -197,6 +215,7 @@ static int parquetColumn(
   sqlite3_context *ctx,       /* First argument to sqlite3_result_...() */
   int col                       /* Which column to return */
 ){
+  try {
     ParquetCursor *cursor = ((sqlite3_vtab_cursor_parquet*)cur)->cursor;
     cursor->ensureColumn(col);
 
@@ -255,6 +274,11 @@ static int parquetColumn(
       }
     }
     return SQLITE_OK;
+  } catch(std::bad_alloc& ba) {
+    return SQLITE_NOMEM;
+  } catch(std::exception& e) {
+    return SQLITE_ERROR;
+  }
 }
 
 /*
@@ -382,6 +406,7 @@ static int parquetFilter(
   int argc,
   sqlite3_value **argv
 ){
+  try {
     ParquetCursor* cursor = ((sqlite3_vtab_cursor_parquet*)cur)->cursor;
     sqlite3_index_info* indexInfo = (sqlite3_index_info*)idxStr;
 
@@ -438,6 +463,11 @@ static int parquetFilter(
     }
     cursor->reset(constraints);
     return parquetNext(cur);
+  } catch(std::bad_alloc& ba) {
+    return SQLITE_NOMEM;
+  } catch(std::exception& e) {
+    return SQLITE_ERROR;
+  }
 }
 
 /*
@@ -450,6 +480,7 @@ static int parquetBestIndex(
   sqlite3_vtab *tab,
   sqlite3_index_info *pIdxInfo
 ){
+  try {
 
 #ifdef DEBUG
     ParquetTable* table = ((sqlite3_vtab_parquet*)tab)->table;
@@ -510,6 +541,11 @@ static int parquetBestIndex(
     }
 
     return SQLITE_OK;
+  } catch(std::bad_alloc& ba) {
+    return SQLITE_NOMEM;
+  } catch(std::exception& e) {
+    return SQLITE_ERROR;
+  }
 }
 
 

tests/test-failmalloc

@@ -3,6 +3,9 @@ set -euo pipefail
 
 # A harness that runs SQLite with the parquet extension in an environment where malloc randomly
 # fails. "Success" is if the logs don't have any C++ exceptions that talk about std::bad_alloc
+#
+# The results can need a bit of interpretation; look at the log and see if it sniffs like
+# the segfault came from Python or SQLite.
 
 ensure_failmalloc() {
   if [ ! -d libfailmalloc ]; then
@@ -19,11 +22,16 @@ ensure_failmalloc() {
 run_under_low_memory() {
   start=$(date +%s%3N)
   set +e
-  env LD_PRELOAD="$here"/libfailmalloc/.libs/libfailmalloc.so FAILMALLOC_PROBABILITY=0.00001 ./test-random &> results.bad_alloc
+  env LD_PRELOAD="$here"/libfailmalloc/.libs/libfailmalloc.so FAILMALLOC_PROBABILITY=0.00001 ./test-random >results.bad_alloc 2>&1
-  set -e
+  rv=$?
   now=$(date +%s%3N)
   echo "Bailed after $((now-start)) ms"
-  ! grep std::bad_alloc results.bad_alloc
+  set -e
+  if [ "$rv" -gt 127 ]; then
+    cat results.bad_alloc
+    echo "Segfaulted with exit code: $rv"
+    exit 1
+  fi
 }
 
 main() {
@@ -33,7 +41,7 @@ main() {
 
   ensure_failmalloc
   # Sometimes we'll exit due to a Python memory issue, so try a few times.
-  for i in {0..10}; do
+  for i in {0..100}; do
     run_under_low_memory
   done
 }