conntrack: merge pull request #695 from ffourcot/ram_conntrack

Enable generator on Conntrack() and reduce memory usage

Bug-Url: https://github.com/svinota/pyroute2/pull/695

pyroute2/conntrack.py

@@ -9,6 +9,9 @@ from pyroute2.netlink.nfnetlink.nfctsocket import NFCTSocket
 
 class NFCTATcpProtoInfo(object):
 
+    __slots__ = ('state', 'wscale_orig', 'wscale_reply', 'flags_orig',
+                 'flags_reply')
+
     def __init__(self, state, wscale_orig=None, wscale_reply=None,
                  flags_orig=None, flags_reply=None):
         self.state = state
@@ -49,6 +52,9 @@ class NFCTATcpProtoInfo(object):
 
 class ConntrackEntry(object):
 
+    __slots__ = ('tuple_orig', 'tuple_reply', 'status', 'timeout',
+                 'protoinfo', 'mark', 'id', 'use')
+
     def __init__(self, family, tuple_orig, tuple_reply, cta_status,
                  cta_timeout, cta_protoinfo, cta_mark, cta_id, cta_use):
         self.tuple_orig = NFCTAttrTuple.from_netlink(family, tuple_orig)
@@ -86,6 +92,8 @@ class Conntrack(NFCTSocket):
     """
     High level conntrack functions
     """
+    def __init__(self, nlm_generator=True, **kwargs):
+        super(Conntrack, self).__init__(nlm_generator=nlm_generator, **kwargs)
 
     def stat(self):
         """ Return current statistics per CPU
@@ -107,16 +115,16 @@ class Conntrack(NFCTSocket):
         Same result than /proc/sys/net/netfilter/nf_conntrack_count file
         or conntrack -C command
         """
-        ndmsg = super(Conntrack, self).count()
-        return ndmsg[0].get_attr('CTA_STATS_GLOBAL_ENTRIES')
+        for ndmsg in super(Conntrack, self).count():
+            return ndmsg.get_attr('CTA_STATS_GLOBAL_ENTRIES')
 
     def conntrack_max_size(self):
         """
         Return the max size of connection tracking table
         /proc/sys/net/netfilter/nf_conntrack_max
         """
-        ndmsg = super(Conntrack, self).conntrack_max_size()
-        return ndmsg[0].get_attr('CTA_STATS_GLOBAL_MAX_ENTRIES')
+        for ndmsg in super(Conntrack, self).conntrack_max_size():
+            return ndmsg.get_attr('CTA_STATS_GLOBAL_MAX_ENTRIES')
 
     def delete(self, entry):
         if isinstance(entry, ConntrackEntry):
@@ -125,7 +133,12 @@ class Conntrack(NFCTSocket):
             tuple_orig = entry
         else:
             raise NotImplementedError()
-        self.entry('del', tuple_orig=tuple_orig)
+        for ndmsg in self.entry('del', tuple_orig=tuple_orig):
+            return ndmsg
+
+    def entry(self, cmd, **kwargs):
+        for res in super(Conntrack, self).entry(cmd, **kwargs):
+            return res
 
     def dump_entries(self, mark=None, mark_mask=0xffffffff, tuple_orig=None,
                      tuple_reply=None):

pyroute2/netlink/nfnetlink/nfctsocket.py

@@ -367,6 +367,11 @@ class NFCTAttr(object):
 
 
 class NFCTAttrTuple(NFCTAttr):
+
+    __slots__ = ('saddr', 'daddr', 'proto', 'sport', 'dport',
+                 'icmp_id', 'icmp_type', 'family', '_attr_ip',
+                 '_attr_icmp')
+
     def __init__(self, family=socket.AF_INET,
                  saddr=None, daddr=None, proto=None, sport=None, dport=None,
                  icmp_id=None, icmp_type=None, icmp_code=None):
@@ -587,8 +592,8 @@ class NFCTSocket(NetlinkSocket):
         IPCTNL_MSG_CT_GET_UNCONFIRMED: nfct_msg,
     }.items())
 
-    def __init__(self, nfgen_family=socket.AF_INET):
-        super(NFCTSocket, self).__init__(family=NETLINK_NETFILTER)
+    def __init__(self, nfgen_family=socket.AF_INET, **kwargs):
+        super(NFCTSocket, self).__init__(family=NETLINK_NETFILTER, **kwargs)
         self.register_policy(self.policy)
         self._nfgen_family = nfgen_family
 

tests/general/test_conntrack.py

@@ -86,7 +86,7 @@ class TestConntrack(BasicSetup):
         # These values should be pretty the same, but the call is not atomic
         # so some sessions may end or begin that time. Thus the difference
         # may occur, but should not be significant.
-        assert abs(len(self.ct.dump()) - self.ct.count()) < 3
+        assert abs(len(list(self.ct.dump())) - self.ct.count()) < 3
 
 
 class TestNFCTSocket(BasicSetup):