Implement reference counting mechanism

tfw/internals/ref_counter/__init__.py

@@ -0,0 +1 @@
+from .ref_counter import RefCounter

tfw/internals/ref_counter/ref_counter.py

@@ -0,0 +1,40 @@
+from os import remove
+from fcntl import flock, LOCK_EX, LOCK_UN
+
+
+class RefCounter:
+    def __init__(self, lockpath):
+        self.lockpath = lockpath
+        self._lockfile = open(self.lockpath, 'a+')
+        flock(self._lockfile, LOCK_EX)
+        counter = self._read_counter()
+        self._write_counter(counter+1)
+        flock(self._lockfile, LOCK_UN)
+
+    def _read_counter(self):
+        self._lockfile.seek(0)
+        try:
+            counter = int(self._lockfile.read())
+        except ValueError:
+            counter = 0
+        return counter
+
+    def _write_counter(self, counter):
+        self._lockfile.seek(0)
+        self._lockfile.truncate()
+        self._lockfile.write(str(counter))
+        self._lockfile.flush()
+
+    def teardown_instance(self):
+        flock(self._lockfile, LOCK_EX)
+        counter = self._read_counter()
+        if counter <= 1:
+            remove(self.lockpath)
+            self.deallocate()
+        else:
+            self._write_counter(counter-1)
+            flock(self._lockfile, LOCK_UN)
+        self._lockfile.close()
+
+    def deallocate(self):
+        pass

tfw/internals/ref_counter/test_ref_counter.py

@@ -0,0 +1,69 @@
+# pylint: disable=redefined-outer-name
+from dataclasses import dataclass
+from textwrap import dedent
+from os import mkfifo
+from os.path import join
+from signal import SIGINT
+from subprocess import DEVNULL, Popen, PIPE
+from tempfile import TemporaryDirectory
+
+import pytest
+
+from .ref_counter import RefCounter
+
+
+@dataclass
+class CounterContext:
+    lockpath: str
+    pipepath: str
+
+    @property
+    def source(self):
+        return dedent(f'''\
+            from time import sleep
+            from atexit import register
+            from ref_counter import RefCounter
+
+            counter = RefCounter('{self.lockpath}')
+            register(counter.teardown_instance)
+            print(flush=True)
+            while True:
+                sleep(1)
+        ''')
+
+
+@pytest.fixture
+def context():
+    with TemporaryDirectory() as workdir:
+        pipepath = join(workdir, 'test.pipe')
+        mkfifo(pipepath)
+        yield CounterContext(join(workdir, 'test.lock'), pipepath)
+
+def test_increment_decrement(context):
+    counter, processes = 0, []
+    for _ in range(5):
+        new_proc = Popen(['python3', '-c', context.source], stdout=PIPE, stderr=DEVNULL)
+        new_proc.stdout.readline()
+        processes.append(new_proc)
+        counter += 1
+    for proc in processes:
+        with open(context.lockpath, 'r') as lock:
+            assert lock.read() == str(counter)
+            counter -= 1
+        proc.send_signal(SIGINT)
+        proc.wait()
+
+def test_deallocate(context):
+    state = False
+    def trigger():
+        nonlocal state
+        state = True
+    refcounters = []
+    for _ in range(32):
+        new_refc = RefCounter(context.lockpath)
+        new_refc.deallocate = trigger
+        refcounters.append(new_refc)
+    for refc in refcounters:
+        assert not state
+        refc.teardown_instance()
+    assert state