191 lines
7.5 KiB
Python
191 lines
7.5 KiB
Python
#!/usr/bin/env python3
|
|
|
|
"""
|
|
This file belongs to https://github.com/bitkeks/python-netflow-v9-softflowd.
|
|
|
|
Copyright 2016-2020 Dominik Pataky <dev@bitkeks.eu>
|
|
Licensed under MIT License. See LICENSE.
|
|
"""
|
|
import io
|
|
import linecache
|
|
import cProfile
|
|
import pathlib
|
|
import pstats
|
|
import resource
|
|
import tracemalloc
|
|
import unittest
|
|
from pstats import SortKey
|
|
|
|
from tests.lib import send_recv_packets, generate_packets
|
|
|
|
NUM_PACKETS_PERFORMANCE = 2000
|
|
|
|
|
|
class TestNetflowIPFIXPerformance(unittest.TestCase):
|
|
def setUp(self) -> None:
|
|
"""
|
|
Before each test run, start tracemalloc profiling.
|
|
:return:
|
|
"""
|
|
tracemalloc.start()
|
|
print("\n\n")
|
|
|
|
def tearDown(self) -> None:
|
|
"""
|
|
After each test run, stop tracemalloc.
|
|
:return:
|
|
"""
|
|
tracemalloc.stop()
|
|
|
|
def _memory_of_version(self, version, store_packets=500) -> tracemalloc.Snapshot:
|
|
"""
|
|
Create memory snapshot of collector run with packets of version :version:
|
|
:param version:
|
|
:return:
|
|
"""
|
|
if not tracemalloc.is_tracing():
|
|
raise RuntimeError
|
|
pkts, t1, t2 = send_recv_packets(generate_packets(NUM_PACKETS_PERFORMANCE, version),
|
|
store_packets=store_packets)
|
|
self.assertEqual(len(pkts), NUM_PACKETS_PERFORMANCE)
|
|
snapshot = tracemalloc.take_snapshot()
|
|
del pkts
|
|
return snapshot
|
|
|
|
@staticmethod
|
|
def _print_memory_statistics(snapshot: tracemalloc.Snapshot, key: str, topx: int = 10):
|
|
"""
|
|
Print memory statistics from a tracemalloc.Snapshot in certain formats.
|
|
:param snapshot:
|
|
:param key:
|
|
:param topx:
|
|
:return:
|
|
"""
|
|
if key not in ["filename", "lineno", "traceback"]:
|
|
raise KeyError
|
|
|
|
stats = snapshot.statistics(key)
|
|
if key == "lineno":
|
|
for idx, stat in enumerate(stats[:topx]):
|
|
frame = stat.traceback[0]
|
|
print("\n{idx:02d}: {filename}:{lineno} {size:.1f} KiB, count {count}".format(
|
|
idx=idx+1, filename=frame.filename, lineno=frame.lineno, size=stat.size / 1024, count=stat.count
|
|
))
|
|
|
|
lines = []
|
|
lines_whitespaces = []
|
|
for lineshift in range(-3, 2):
|
|
stat = linecache.getline(frame.filename, frame.lineno + lineshift)
|
|
lines_whitespaces.append(len(stat) - len(stat.lstrip(" "))) # count
|
|
lines.append(stat.strip())
|
|
lines_whitespaces = [x - min([y for y in lines_whitespaces if y > 0]) for x in lines_whitespaces]
|
|
for lidx, stat in enumerate(lines):
|
|
print(" {}{}".format("> " if lidx == 3 else "| ", " " * lines_whitespaces.pop(0) + stat))
|
|
elif key == "filename":
|
|
for idx, stat in enumerate(stats[:topx]):
|
|
frame = stat.traceback[0]
|
|
print("{idx:02d}: {filename:80s} {size:6.1f} KiB, count {count:5<d}".format(
|
|
idx=idx + 1, filename=frame.filename, size=stat.size / 1024, count=stat.count
|
|
))
|
|
|
|
def test_compare_memory(self):
|
|
"""
|
|
Test memory usage of two collector runs with IPFIX and NetFlow v9 packets respectively.
|
|
Then compare the two memory snapshots to make sure the libraries do not cross each other.
|
|
TODO: more features could be tested, e.g. too big of a difference if one version is optimized better
|
|
:return:
|
|
"""
|
|
pkts, t1, t2 = send_recv_packets(generate_packets(NUM_PACKETS_PERFORMANCE, 10))
|
|
self.assertEqual(len(pkts), NUM_PACKETS_PERFORMANCE)
|
|
snapshot_ipfix = tracemalloc.take_snapshot()
|
|
del pkts
|
|
tracemalloc.clear_traces()
|
|
|
|
pkts, t1, t2 = send_recv_packets(generate_packets(NUM_PACKETS_PERFORMANCE, 9))
|
|
self.assertEqual(len(pkts), NUM_PACKETS_PERFORMANCE)
|
|
snapshot_v9 = tracemalloc.take_snapshot()
|
|
del pkts
|
|
|
|
stats = snapshot_v9.compare_to(snapshot_ipfix, "lineno")
|
|
for stat in stats:
|
|
if stat.traceback[0].filename.endswith("netflow/ipfix.py"):
|
|
self.assertEqual(stat.count, 0)
|
|
self.assertEqual(stat.size, 0)
|
|
|
|
stats = snapshot_ipfix.compare_to(snapshot_v9, "lineno")
|
|
for stat in stats:
|
|
if stat.traceback[0].filename.endswith("netflow/v9.py"):
|
|
self.assertEqual(stat.count, 0)
|
|
self.assertEqual(stat.size, 0)
|
|
|
|
def test_memory_ipfix(self):
|
|
"""
|
|
Test memory usage of the collector with IPFIX packets.
|
|
Three iterations are done with different amounts of packets to be stored.
|
|
With this approach, increased usage of memory can be captured when the ExportPacket objects are not deleted.
|
|
:return:
|
|
"""
|
|
for store_pkts in [0, 500, -1]: # -1 is compatibility value for "store all"
|
|
snapshot_ipfix = self._memory_of_version(10, store_packets=store_pkts)
|
|
|
|
# TODO: this seems misleading, maybe reads the memory of the whole testing process?
|
|
# system_memory = pathlib.Path("/proc/self/statm").read_text()
|
|
# pagesize = resource.getpagesize()
|
|
# print("Total RSS memory used: {:.1f} KiB".format(int(system_memory.split()[1]) * pagesize // 1024.))
|
|
|
|
print("\nIPFIX memory usage with {} packets being stored".format(store_pkts))
|
|
self._print_memory_statistics(snapshot_ipfix, "filename")
|
|
if store_pkts == -1:
|
|
# very verbose and most interesting in iteration with all ExportPackets being stored in memory
|
|
print("\nTraceback for run with all packets being stored in memory")
|
|
self._print_memory_statistics(snapshot_ipfix, "lineno")
|
|
|
|
def test_memory_v1(self):
|
|
"""
|
|
Test memory with NetFlow v1
|
|
:return:
|
|
"""
|
|
snapshot_v1 = self._memory_of_version(1)
|
|
print("\nNetFlow v1 memory usage by file")
|
|
self._print_memory_statistics(snapshot_v1, "filename")
|
|
|
|
def test_memory_v5(self):
|
|
"""
|
|
Test memory with NetFlow v5
|
|
:return:
|
|
"""
|
|
snapshot_v5 = self._memory_of_version(5)
|
|
print("\nNetFlow v5 memory usage by file")
|
|
self._print_memory_statistics(snapshot_v5, "filename")
|
|
|
|
def test_memory_v9(self):
|
|
"""
|
|
Test memory usage of the collector with NetFlow v9 packets.
|
|
:return:
|
|
"""
|
|
snapshot_v9 = self._memory_of_version(9)
|
|
print("\nNetFlow v9 memory usage by file")
|
|
self._print_memory_statistics(snapshot_v9, "filename")
|
|
print("\nNetFlow v9 memory usage by line")
|
|
self._print_memory_statistics(snapshot_v9, "lineno")
|
|
|
|
@unittest.skip("Does not work as expected due to threading")
|
|
def test_time_ipfix(self):
|
|
"""
|
|
Profile function calls and CPU time.
|
|
TODO: this does not work with threading in the collector, yet
|
|
:return:
|
|
"""
|
|
profile = cProfile.Profile()
|
|
profile.enable(subcalls=True, builtins=True)
|
|
pkts, t1, t2 = send_recv_packets(generate_packets(NUM_PACKETS_PERFORMANCE, 10), delay=0, store_packets=500)
|
|
self.assertEqual(len(pkts), NUM_PACKETS_PERFORMANCE)
|
|
profile.disable()
|
|
|
|
for sort_by in [SortKey.CUMULATIVE, SortKey.CALLS]:
|
|
s = io.StringIO()
|
|
ps = pstats.Stats(profile, stream=s)
|
|
ps.sort_stats(sort_by).print_stats("netflow")
|
|
ps.sort_stats(sort_by).print_callees(.5)
|
|
print(s.getvalue())
|