Solve performance problems (disable WAL on DB)

Author: grafolean

dbutils.py

@@ -135,13 +135,22 @@ def migration_step_1():
 
 def migration_step_2():
     with db.cursor() as c:
-        c.execute(f'CREATE TABLE {DB_PREFIX}records (seq BIGSERIAL NOT NULL PRIMARY KEY, ts TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, client_ip TEXT, version SMALLSERIAL NOT NULL);')
-        c.execute(f'CREATE TABLE {DB_PREFIX}flows (record INTEGER NOT NULL REFERENCES {DB_PREFIX}records(seq) ON DELETE CASCADE, data JSONB NOT NULL);')
-
-def migration_step_3():
-    with db.cursor() as c:
+        # UNLOGGED: Disabling WAL avoids high I/O load. Since NetFlow data is of temporary nature, this still
+        # allows us to perform queries, but if the database crashes it is acceptable to lose all of the records.
+        c.execute(f'CREATE UNLOGGED TABLE {DB_PREFIX}records (seq BIGSERIAL NOT NULL PRIMARY KEY, ts TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, client_ip TEXT);')
         c.execute(f'CREATE INDEX {DB_PREFIX}records_ts ON {DB_PREFIX}records (ts);')
 
-def migration_step_4():
-    with db.cursor() as c:
-        c.execute(f'ALTER TABLE {DB_PREFIX}records DROP COLUMN version;')
+        c.execute(f"""
+            CREATE UNLOGGED TABLE {DB_PREFIX}flows (
+                record INTEGER NOT NULL REFERENCES {DB_PREFIX}records(seq) ON DELETE CASCADE,
+                IN_BYTES INTEGER,
+                PROTOCOL SMALLINT,
+                DIRECTION SMALLINT,
+                L4_DST_PORT INTEGER,
+                L4_SRC_PORT INTEGER,
+                INPUT_SNMP SMALLINT,
+                OUTPUT_SNMP SMALLINT,
+                IPV4_DST_ADDR TEXT,
+                IPV4_SRC_ADDR TEXT
+            );
+        """)

netflowbot.py

@@ -117,18 +117,18 @@ def get_values_traffic_in(output_path_prefix, from_time, to_time):
             # TODO: missing check for IP: r.client_ip = %s AND
             c.execute(f"""
                 SELECT
-                    f.data->'INPUT_SNMP',
-                    sum((f.data->'IN_BYTES')::integer)
+                    f.INPUT_SNMP,
+                    sum(f.IN_BYTES)
                 FROM
                     {DB_PREFIX}records "r",
                     {DB_PREFIX}flows "f"
                 WHERE
                     r.ts >= %s AND
                     r.ts < %s AND
                     r.seq = f.record AND
-                    ((f.data->'DIRECTION')::integer) = 1
+                    f.DIRECTION = 0
                 GROUP BY
-                    f.data->'INPUT_SNMP'
+                    f.INPUT_SNMP
             """, (from_time, to_time,))
 
             values = []
@@ -146,18 +146,18 @@ def get_values_traffic_out(output_path_prefix, from_time, to_time):
             # TODO: missing check for IP: r.client_ip = %s AND
             c.execute(f"""
                 SELECT
-                    f.data->'OUTPUT_SNMP',
-                    sum((f.data->'IN_BYTES')::integer)
+                    f.OUTPUT_SNMP,
+                    sum(f.IN_BYTES)
                 FROM
                     {DB_PREFIX}records "r",
                     {DB_PREFIX}flows "f"
                 WHERE
                     r.ts >= %s AND
                     r.ts < %s AND
                     r.seq = f.record AND
-                    ((f.data->'DIRECTION')::integer) = 1
+                    f.DIRECTION = 1
                 GROUP BY
-                    f.data->'OUTPUT_SNMP'
+                    f.OUTPUT_SNMP
             """, (from_time, to_time,))
 
             values = []
@@ -175,19 +175,19 @@ def get_top_N_IPs(output_path_prefix, from_time, to_time, interface_index, is_di
             # TODO: missing check for IP: r.client_ip = %s AND
             c.execute(f"""
                 SELECT
-                    f.data->'IPV4_DST_ADDR',
-                    sum((f.data->'IN_BYTES')::integer) "traffic"
+                    f.IPV4_DST_ADDR,
+                    sum(f.IN_BYTES) "traffic"
                 FROM
                     netflow_records "r",
                     netflow_flows "f"
                 WHERE
                     r.ts >= %s AND
                     r.ts < %s AND
                     r.seq = f.record AND
-                    (f.data->'{'INPUT_SNMP' if is_direction_in else 'OUTPUT_SNMP'}')::integer = %s AND
-                    (f.data->'DIRECTION')::integer = {'0' if is_direction_in else '1'}
+                    f.{'INPUT_SNMP' if is_direction_in else 'OUTPUT_SNMP'} = %s AND
+                    f.DIRECTION = {'0' if is_direction_in else '1'}
                 GROUP BY
-                    f.data->'IPV4_DST_ADDR'
+                    f.IPV4_DST_ADDR
                 ORDER BY
                     traffic desc
                 LIMIT 10;

netflowwriter.py

@@ -54,17 +54,34 @@ def write_record(j):
         # then save each of the flows within the record, but use execute_values() to perform bulk insert:
         def _get_data(record_seq, flows):
             for flow in flows:
-                yield (record_seq, flow,)
+                yield (
+                    record_seq,
+                    flow.get('IN_BYTES'),
+                    flow.get('PROTOCOL'),
+                    flow.get('DIRECTION'),
+                    flow.get('L4_DST_PORT'),
+                    flow.get('L4_SRC_PORT'),
+                    flow.get('INPUT_SNMP'),
+                    flow.get('OUTPUT_SNMP'),
+                    flow.get('IPV4_DST_ADDR'),
+                    flow.get('IPV4_SRC_ADDR'),
+                )
         data_iterator = _get_data(record_seq, j['flows'])
-        psycopg2.extras.execute_values(c, f"INSERT INTO {DB_PREFIX}flows (record, data) VALUES %s", data_iterator, "(%s, %s)", page_size=100)
-
+        psycopg2.extras.execute_values(
+            c,
+            f"INSERT INTO {DB_PREFIX}flows (record, IN_BYTES, PROTOCOL, DIRECTION, L4_DST_PORT, L4_SRC_PORT, INPUT_SNMP, OUTPUT_SNMP, IPV4_DST_ADDR, IPV4_SRC_ADDR) VALUES %s",
+            data_iterator,
+            "(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)",
+            page_size=100
+        )
 
 if __name__ == "__main__":
     NAMED_PIPE_FILENAME = os.environ.get('NAMED_PIPE_FILENAME', None)
     if not NAMED_PIPE_FILENAME:
         raise Exception("Please specify NAMED_PIPE_FILENAME environment var")
 
     migrate_if_needed()
+
     try:
         process_named_pipe(NAMED_PIPE_FILENAME)
     except KeyboardInterrupt: