SNOW-1015703: fix mem leak and fix flaw in mem leak detect program (#1867)

Author: sfc-gh-aling

DESCRIPTION.md

@@ -8,6 +8,9 @@ Source code is also available at: https://github.com/snowflakedb/snowflake-conne
 
 # Release Notes
 
+- v3.7.1(TBD)
+  - Fixed a memory leak in decimal data conversion.
+
 - v3.7.0(January 25,2024)
 
   - Added a new boolean parameter `force_return_table` to `SnowflakeCursor.fetch_arrow_all` to force returning `pyarrow.Table` in case of zero rows.

src/snowflake/connector/nanoarrow_cpp/ArrowIterator/DecimalConverter.cpp

@@ -84,8 +84,18 @@ PyObject* DecimalFromDecimalConverter::toPyObject(int64_t rowIndex) const {
   ArrowDecimalGetBytes(&arrowDecimal, outBytes);
   PyObject* int128_bytes = PyBytes_FromStringAndSize(&outBytes, 16);
   */
-  return PyObject_CallMethod(m_context, "DECIMAL128_to_decimal", "Si",
-                             int128_bytes, m_scale);
+  PyObject* return_object = PyObject_CallMethod(
+      m_context, "DECIMAL128_to_decimal", "Si", int128_bytes, m_scale);
+  /**
+  int128_bytes is a new referenced created by PyBytes_FromStringAndSize,
+  to avoid memory leak we need to free it after usage
+  check docs:
+     https://docs.python.org/3/c-api/bytes.html#c.PyBytes_FromStringAndSize
+     https://docs.python.org/3/c-api/refcounting.html#c.Py_XDECREF
+
+  */
+  Py_XDECREF(int128_bytes);
+  return return_object;
 }
 
 }  // namespace sf

test/stress/README.md

@@ -0,0 +1,30 @@
+## quick start for performance testing
+
+
+### setup
+
+note: you need to put your own credentials into parameters.py
+
+```bash
+git clone [email protected]:snowflakedb/snowflake-connector-python.git
+cd snowflake-connector-python/test/stress
+pip install -r dev_requirements.txt
+touch parameters.py  # set your own connection parameters
+```
+
+### run unit perf test
+
+This test will use the test dataset stored in the "stress_test_data" folder.
+check the read me in the folder to see what datasets are available.
+
+```python
+python local_iterator.py
+```
+
+### run e2e perf test
+
+This test will run query against snowflake. update the script to prepare the data and run the test.
+
+```python
+python e2e_iterator.py
+```

test/stress/e2e_iterator.py

@@ -15,7 +15,7 @@
 import argparse
 
 import util as stress_util
-from util import task_memory_decorator, task_time_execution_decorator
+from util import task_execution_decorator
 
 import snowflake.connector
 from parameters import CONNECTION_PARAMETERS
@@ -124,23 +124,31 @@ def execute_task(task, cursor, table_name, iteration_cnt):
     ) as conn, conn.cursor() as cursor:
         test_table_name = args.test_table_name
 
-        memory_check_task = task_memory_decorator(task_fetch_arrow_batches)
-        execute_task(memory_check_task, cursor, test_table_name, args.iteration_cnt)
-        memory_records = stress_util.collect_memory_records()
-
-        perf_check_task = task_time_execution_decorator(task_fetch_arrow_batches)
-        execute_task(perf_check_task, cursor, test_table_name, args.iteration_cnt)
-        time_records = stress_util.collect_time_execution_records()
-
-        print("average time is", sum(time_records) / len(time_records))
+        perf_record_file = "stress_perf_record"
+        memory_record_file = "stress_memory_record"
+        with open(perf_record_file, "w") as perf_file, open(
+            memory_record_file, "w"
+        ) as memory_file:
+            task = task_execution_decorator(
+                task_fetch_arrow_batches, perf_file, memory_file
+            )
+            execute_task(task, cursor, test_table_name, args.iteration_cnt)
 
         if can_draw:
-            plt.plot([i for i in range(len(time_records))], time_records)
-            plt.title("per iteration execution time")
-            plt.show()
-            plt.plot(
-                [item[0] for item in memory_records],
-                [item[1] for item in memory_records],
-            )
-            plt.title("memory usage")
-            plt.show()
+            with open(perf_record_file) as perf_file, open(
+                memory_record_file
+            ) as memory_file:
+                # sample rate
+                perf_lines = perf_file.readlines()
+                perf_records = [float(line) for line in perf_lines]
+
+                memory_lines = memory_file.readlines()
+                memory_records = [float(line) for line in memory_lines]
+
+                plt.plot([i for i in range(len(perf_records))], perf_records)
+                plt.title("per iteration execution time")
+                plt.show(block=False)
+                plt.figure()
+                plt.plot([i for i in range(len(memory_records))], memory_records)
+                plt.title("memory usage")
+                plt.show(block=True)

test/stress/local_iterator.py

@@ -19,7 +19,7 @@
 import secrets
 
 import util as stress_util
-from util import task_memory_decorator, task_time_execution_decorator
+from util import task_execution_decorator
 
 from snowflake.connector.arrow_context import ArrowConverterContext
 from snowflake.connector.nanoarrow_arrow_iterator import (
@@ -130,13 +130,13 @@ def execute_task(
     parser.add_argument(
         "--iteration_cnt",
         type=int,
-        default=10,
+        default=100000,
         help="how many times to run the test function, default is 100000",
     )
     parser.add_argument(
         "--data_file",
         type=str,
-        default="test_data",
+        default="stress_test_data/test_multi_column_row_decimal_data",
         help="a local file to read data from, the file contains base64 encoded string returned from snowflake",
     )
     parser.add_argument(
@@ -153,53 +153,59 @@ def execute_task(
 
     args = parser.parse_args()
 
-    with open(args.data_file) as f:
-        b64data = f.read()
+    try:
+        # file contains base64 encoded data
+        with open(args.data_file) as f:
+            b64data = f.read()
 
-    decode_bytes = base64.b64decode(b64data)
+        decode_bytes = base64.b64decode(b64data)
+    except UnicodeDecodeError:
+        # file contains raw bytes data
+        with open(args.data_file, "rb") as f:
+            decode_bytes = f.read()
 
     # if connector is pre-release, then it's nanoarrow based iterator
     print(
         "Testing connector version: ",
         ".".join([str(v) for v in VERSION if v is not None]),
     )
 
-    perf_check_task_for_loop_iterator = task_time_execution_decorator(
-        task_for_loop_iterator_expected_error
-        if args.test_error_method
-        else task_for_loop_iterator
-    )
-    memory_check_task_for_loop_iterator = task_memory_decorator(
-        task_for_loop_iterator_expected_error
-        if args.test_error_method
-        else task_for_loop_iterator
-    )
-
-    execute_task(
-        memory_check_task_for_loop_iterator,
-        decode_bytes,
-        create_nanoarrow_pyarrow_iterator,
-        args.iteration_cnt,
-        args.use_table_unit,
-    )
-    memory_records = stress_util.collect_memory_records()
-    execute_task(
-        perf_check_task_for_loop_iterator,
-        decode_bytes,
-        create_nanoarrow_pyarrow_iterator,
-        args.iteration_cnt,
-        args.use_table_unit,
-    )
-    time_records = stress_util.collect_time_execution_records()
+    perf_record_file = "stress_perf_record"
+    memory_record_file = "stress_memory_record"
+    with open(perf_record_file, "w") as perf_file, open(
+        memory_record_file, "w"
+    ) as memory_file:
+        task_for_loop_iterator = task_execution_decorator(
+            task_for_loop_iterator_expected_error
+            if args.test_error_method
+            else task_for_loop_iterator,
+            perf_file,
+            memory_file,
+        )
 
-    print("average time is", sum(time_records) / len(time_records))
+        execute_task(
+            task_for_loop_iterator,
+            decode_bytes,
+            create_nanoarrow_pyarrow_iterator,
+            args.iteration_cnt,
+            args.use_table_unit,
+        )
 
     if can_draw:
-        plt.plot([i for i in range(len(time_records))], time_records)
-        plt.title("per iteration execution time")
-        plt.show()
-        plt.plot(
-            [item[0] for item in memory_records], [item[1] for item in memory_records]
-        )
-        plt.title("memory usage")
-        plt.show()
+        with open(perf_record_file) as perf_file, open(
+            memory_record_file
+        ) as memory_file:
+            # sample rate
+            perf_lines = perf_file.readlines()
+            perf_records = [float(line) for line in perf_lines]
+
+            memory_lines = memory_file.readlines()
+            memory_records = [float(line) for line in memory_lines]
+
+            plt.plot([i for i in range(len(perf_records))], perf_records)
+            plt.title("per iteration execution time")
+            plt.show(block=False)
+            plt.figure()
+            plt.plot([i for i in range(len(memory_records))], memory_records)
+            plt.title("memory usage")
+            plt.show(block=True)

test/stress/stress_test_data/README.md

@@ -0,0 +1,68 @@
+# Test Data Description
+
+## test_data_all_types
+
+This dataset contains 1 batch, 100 rows of data,  and the schema of the data contains 27 columns.
+
+### sample row data:
+
+```python
+(
+    123456,
+    bytearray(b'HELP'),
+    True,
+    'a',
+    'b',
+    datetime.date(2023, 7, 18),
+    datetime.datetime(2023, 7, 18, 12, 51),
+    Decimal('984.280'),
+    Decimal('268.350'),
+    123.456,
+    738.132,
+    6789,
+    23456,
+    12583,
+    513.431,
+    10,
+    9,
+    'fjisfsj',
+    'wkdoajde131',
+    datetime.time(12, 34, 56),
+    datetime.datetime(2021, 1, 1, 0, 0),
+    datetime.datetime(2021, 1, 1, 0, 0, tzinfo=<UTC>),
+    datetime.datetime(2020, 12, 31, 16, 0, tzinfo=<DstTzInfo 'America/Los_Angeles' PST-1 day, 16:00:00 STD>),
+    datetime.datetime(2021, 1, 1, 0, 0),
+    1,
+    bytearray(b'HELP'),
+    'vxlmls!21321#@!#!'
+)
+```
+
+## test_multi_column_row_decimal_data
+
+This dataset contains 9 batches, each batch has approximately ~1700 rows of data, and the schema of the data contains 19 columns.
+
+### sample row data:
+```python
+(
+    datetime.date(2021, 1, 3),
+    8371,
+    'segment_no_0',
+    1,
+    7,
+    2,
+    Decimal('0.285714'),
+    Decimal('1.000'),
+    Decimal('7.000'),
+    Decimal('2.000'),
+    Decimal('0.285714000'),
+    Decimal('1.000'),
+    Decimal('7.000'),
+    Decimal('2.000'),
+    Decimal('0.285714000'),
+    Decimal('1.000'),
+    Decimal('7.000'),
+    Decimal('2.000'),
+    Decimal('0.285714000')
+)
+```

test/stress/test_data renamed to test/stress/stress_test_data/test_data_all_types

@@ -8,50 +8,26 @@
 
 process = psutil.Process()
 
-last_memory_record = None
-memory_records = []
-time_records = []
-memory_decoration_execution_time = 0
-print_to_console = False
+SAMPLE_RATE = 10  # record data evey SAMPLE_RATE execution
 
 
-def collect_memory_records():
-    memory_records.append(last_memory_record)
-    return memory_records
-
-
-def collect_time_execution_records():
-    return time_records
-
-
-def task_memory_decorator(func):
-    memory_records.clear()
-    global memory_decoration_execution_time
-    memory_decoration_execution_time = 0
-
-    def wrapper(*args, **kwargs):
-        global memory_decoration_execution_time
-        global print_to_console
-        global last_memory_record
-        func(*args, **kwargs)
-        percent = process.memory_percent()
-        if not memory_records or (memory_records and percent != memory_records[-1][1]):
-            memory_records.append((memory_decoration_execution_time, percent))
-        memory_decoration_execution_time += 1
-        last_memory_record = (memory_decoration_execution_time, percent)
-        if print_to_console:
-            print(memory_decoration_execution_time, percent)
-
-    return wrapper
-
-
-def task_time_execution_decorator(func):
-    time_records.clear()
+def task_execution_decorator(func, perf_file, memory_file):
+    count = 0
 
     def wrapper(*args, **kwargs):
         start = time.time()
         func(*args, **kwargs)
+        memory_usage = (
+            process.memory_info().rss / 1024 / 1024
+        )  # rss is of unit bytes, we get unit in MB
         period = time.time() - start
-        time_records.append(period)
+        nonlocal count
+        if count % SAMPLE_RATE == 0:
+            perf_file.write(str(period) + "\n")
+            print(f"execution time {count}")
+            print(f"memory usage: {memory_usage} MB")
+            print(f"execution time: {period} s")
+            memory_file.write(str(memory_usage) + "\n")
+        count += 1
 
     return wrapper