dsmq 1.0.0__py3-none-any.whl → 1.1.0__py3-none-any.whl

dsmq/client.py

@@ -1,9 +1,13 @@
 import json
+import time
 from websockets.sync.client import connect as ws_connect
 
 _default_host = "127.0.0.1"
 _default_port = 30008
 
+_n_retries = 10
+_initial_retry = 0.01  # seconds
+
 
 def connect(host=_default_host, port=_default_port):
     return DSMQClientSideConnection(host, port)
@@ -12,16 +16,50 @@ def connect(host=_default_host, port=_default_port):
 class DSMQClientSideConnection:
     def __init__(self, host, port):
         self.uri = f"ws://{host}:{port}"
-        self.websocket = ws_connect(self.uri)
+        print(f"Connecting to dsmq server at {self.uri}")
+        for i_retry in range(_n_retries):
+            try:
+                self.websocket = ws_connect(self.uri)
+                break
+            except ConnectionRefusedError:
+                self.websocket = None
+                # Exponential backoff
+                # Wait twice as long each time before trying again.
+                time.sleep(_initial_retry * 2**i_retry)
+                print("    ...trying again")
+
+        if self.websocket is None:
+            raise ConnectionRefusedError("Could not connect to dsmq server.")
+
+        self.time_of_last_request = time.time()
 
     def get(self, topic):
         msg = {"action": "get", "topic": topic}
         self.websocket.send(json.dumps(msg))
-
         msg_text = self.websocket.recv()
         msg = json.loads(msg_text)
         return msg["message"]
 
+    def get_wait(self, topic):
+        """
+        A variant of `get()` that retries a few times until it gets
+        a non-empty message. Adjust `_n_tries` and `_initial_retry`
+        to change how persistent it will be.
+        """
+        for i_retry in range(_n_retries):
+            message = self.get(topic)
+            if message != "":
+                return message
+            time.sleep(_initial_retry * 2**i_retry)
+        return message
+
     def put(self, topic, msg_body):
         msg_dict = {"action": "put", "topic": topic, "message": msg_body}
         self.websocket.send(json.dumps(msg_dict))
+
+    def shutdown_server(self):
+        msg_dict = {"action": "shutdown", "topic": ""}
+        self.websocket.send(json.dumps(msg_dict))
+
+    def close(self):
+        self.websocket.close()

dsmq/{demo_linux.py → demo.py}

@@ -1,11 +1,8 @@
-# Heads up: This script only works on Linux because of use of "fork" method.
 import multiprocessing as mp
 from dsmq.server import serve
 import dsmq.example_get_client
 import dsmq.example_put_client
 
-mp.set_start_method("fork")
-
 HOST = "127.0.0.1"
 PORT = 25252
 

dsmq/server.py

@@ -1,6 +1,8 @@
 import json
+import os
 import sqlite3
 import sys
+from threading import Thread
 import time
 from websockets.sync.server import serve as ws_serve
 
@@ -8,30 +10,57 @@ _default_host = "127.0.0.1"
 _default_port = 30008
 _n_retries = 5
 _first_retry = 0.01  # seconds
+_pause = 0.01  # seconds
 _time_to_live = 600.0  # seconds
 
+_db_name = "file::memory:?cache=shared"
+
+# Make this global so it's easy to share
+dsmq_server = None
+
 
 def serve(host=_default_host, port=_default_port):
     """
     For best results, start this running in its own process and walk away.
     """
-    sqlite_conn = sqlite3.connect("file:mem1?mode=memory&cache=shared")
+    # Cleanup temp files.
+    # Under some condition
+    # (which I haven't yet been able to pin down)
+    # a file is generated with the db name.
+    # If it is not removed, it gets
+    # treated as a SQLite db on disk,
+    # which dramatically slows it down,
+    # especially the way it's used here for
+    # rapid-fire one-item reads and writes.
+    filenames = os.listdir()
+    for filename in filenames:
+        if filename[: len(_db_name)] == _db_name:
+            os.remove(filename)
+
+    sqlite_conn = sqlite3.connect(_db_name)
     cursor = sqlite_conn.cursor()
     cursor.execute("""
 CREATE TABLE IF NOT EXISTS messages (timestamp DOUBLE, topic TEXT, message TEXT)
     """)
 
-    with ws_serve(request_handler, host, port) as server:
-        server.serve_forever()
-        print()
-        print(f"Server started at {host} on port {port}.")
-        print("Waiting for clients...")
+    # Making this global in scope is a way to make it available
+    # to the shutdown operation. It's an awkward construction,
+    # and a method of last resort. (If you stumble across this and
+    # figure out something more elegant, please submit a PR!
+    # or send it to me at brohrer@gmail.com,
+    global dsmq_server
+
+    dsmq_server = ws_serve(request_handler, host, port)
+    dsmq_server.serve_forever()
+    print()
+    print(f"Server started at {host} on port {port}.")
+    print("Waiting for clients...")
 
     sqlite_conn.close()
 
 
 def request_handler(websocket):
-    sqlite_conn = sqlite3.connect("file:mem1?mode=memory&cache=shared")
+    sqlite_conn = sqlite3.connect(_db_name)
     cursor = sqlite_conn.cursor()
 
     client_creation_time = time.time()
@@ -85,7 +114,7 @@ FROM messages,
 SELECT MIN(timestamp) AS min_time
 FROM messages
 WHERE topic = :topic
-    AND timestamp > :last_read_time
+AND timestamp > :last_read_time
 ) a
 WHERE topic = :topic
 AND timestamp = a.min_time
@@ -108,9 +137,27 @@ AND timestamp = a.min_time
                 message = ""
 
             websocket.send(json.dumps({"message": message}))
+        elif msg["action"] == "shutdown":
+            # Run this from a separate thread to prevent deadlock
+            global dsmq_server
+            print("Shutting down the dsmq server.")
+
+            def shutdown_gracefully(server_to_shutdown):
+                server_to_shutdown.shutdown()
+                time.sleep(_pause)
+
+                filenames = os.listdir()
+                for filename in filenames:
+                    if filename[: len(_db_name)] == _db_name:
+                        try:
+                            os.remove(filename)
+                        except FileNotFoundError:
+                            pass
+
+            Thread(target=shutdown_gracefully, args=(dsmq_server,)).start()
+            sqlite_conn.close()
         else:
-            print("Action must either be 'put' or 'get'")
-
+            print("Action must either be 'put', 'get', or 'shudown'")
 
         # Periodically clean out messages from the queue that are
         # past their sell buy date.
@@ -121,7 +168,7 @@ AND timestamp = a.min_time
 DELETE FROM messages
 WHERE timestamp < :time_threshold
                 """,
-                {"time_threshold": time_of_last_purge}
+                {"time_threshold": time_of_last_purge},
             )
             sqlite_conn.commit()
             time_of_last_purge = time.time()

dsmq/tests/integration_test.py

@@ -0,0 +1,238 @@
+import multiprocessing as mp
+import time
+from dsmq.server import serve
+from dsmq.client import connect
+
+# spawn is the default method on macOS
+# mp.set_start_method('spawn')
+
+host = "127.0.0.1"
+port = 30303
+
+_short_pause = 0.001
+_pause = 0.01
+_long_pause = 0.1
+_very_long_pause = 0.1
+
+
+def test_client_server():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    mq = connect(host, port)
+    read_completes = False
+    write_completes = False
+
+    n_iter = 11
+    for i in range(n_iter):
+        mq.put("test", f"msg_{i}")
+        write_completes = True
+
+    for i in range(n_iter):
+        msg = mq.get("test")
+        read_completes = True
+
+    assert msg
+    assert write_completes
+    assert read_completes
+
+    mq.shutdown_server()
+    mq.close()
+
+    # It takes a sec to shut down the server
+    time.sleep(_long_pause)
+    assert not p_server.is_alive()
+
+
+def test_write_one_read_one():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    write_client = connect(host, port)
+    read_client = connect(host, port)
+
+    write_client.put("test", "test_msg")
+
+    # It takes a moment for the write to complete
+    time.sleep(_pause)
+    msg = read_client.get("test")
+
+    assert msg == "test_msg"
+
+    write_client.shutdown_server()
+    write_client.close()
+    read_client.close()
+
+
+def test_get_wait():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    write_client = connect(host, port)
+    read_client = connect(host, port)
+
+    write_client.put("test", "test_msg")
+
+    msg = read_client.get_wait("test")
+
+    assert msg == "test_msg"
+
+    write_client.shutdown_server()
+    write_client.close()
+    read_client.close()
+
+
+def test_multitopics():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    write_client = connect(host, port)
+    read_client = connect(host, port)
+
+    write_client.put("test_A", "test_msg_A")
+    write_client.put("test_B", "test_msg_B")
+
+    msg_A = read_client.get_wait("test_A")
+    msg_B = read_client.get_wait("test_B")
+
+    assert msg_A == "test_msg_A"
+    assert msg_B == "test_msg_B"
+
+    write_client.shutdown_server()
+    write_client.close()
+    read_client.close()
+
+
+def test_client_history_cutoff():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    write_client = connect(host, port)
+    write_client.put("test", "test_msg")
+    time.sleep(_pause)
+
+    read_client = connect(host, port)
+    msg = read_client.get("test")
+
+    assert msg == ""
+
+    write_client.shutdown_server()
+    write_client.close()
+    read_client.close()
+
+
+def test_two_write_clients():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    write_client_A = connect(host, port)
+    write_client_B = connect(host, port)
+    read_client = connect(host, port)
+
+    write_client_A.put("test", "test_msg_A")
+    # Wait briefly, to ensure the order of writes
+    time.sleep(_pause)
+    write_client_B.put("test", "test_msg_B")
+    msg_A = read_client.get_wait("test")
+    msg_B = read_client.get_wait("test")
+
+    assert msg_A == "test_msg_A"
+    assert msg_B == "test_msg_B"
+
+    write_client_A.shutdown_server()
+    write_client_A.close()
+    write_client_B.close()
+    read_client.close()
+
+
+def test_two_read_clients():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    write_client = connect(host, port)
+    read_client_A = connect(host, port)
+    read_client_B = connect(host, port)
+
+    write_client.put("test", "test_msg")
+    msg_A = read_client_A.get_wait("test")
+    msg_B = read_client_B.get_wait("test")
+
+    assert msg_A == "test_msg"
+    assert msg_B == "test_msg"
+
+    write_client.shutdown_server()
+    write_client.close()
+    read_client_A.close()
+    read_client_B.close()
+
+
+def speed_write(stop_flag):
+    fast_write_client = connect(host, port)
+    while True:
+        fast_write_client.put("speed_test", "speed")
+        if stop_flag.is_set():
+            break
+    fast_write_client.close()
+
+
+def test_speed_writing():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    write_client = connect(host, port)
+    read_client = connect(host, port)
+
+    stop_flag = mp.Event()
+    p_speed_write = mp.Process(target=speed_write, args=(stop_flag,))
+    p_speed_write.start()
+    time.sleep(_pause)
+
+    # time_a = time.time()
+    write_client.put("test", "test_msg")
+    # time_b = time.time()
+    msg = read_client.get_wait("test")
+    # time_c = time.time()
+
+    # write_time = int((time_b - time_a) * 1e6)
+    # read_time = int((time_c - time_b) * 1e6)
+    # print(f"write time: {write_time} us,  read time: {read_time} us")
+
+    assert msg == "test_msg"
+
+    stop_flag.set()
+    p_speed_write.join()
+    read_client.close()
+    write_client.shutdown_server()
+    write_client.close()
+
+
+def speed_read(stop_flag):
+    fast_read_client = connect(host, port)
+    while True:
+        fast_read_client.get("speed_test")
+        if stop_flag.is_set():
+            break
+    fast_read_client.close()
+
+
+def test_speed_reading():
+    p_server = mp.Process(target=serve, args=(host, port))
+    p_server.start()
+    write_client = connect(host, port)
+    read_client = connect(host, port)
+
+    stop_flag = mp.Event()
+    p_speed_read = mp.Process(target=speed_read, args=(stop_flag,))
+    p_speed_read.start()
+
+    # time_a = time.time()
+    write_client.put("test", "test_msg")
+    # time_b = time.time()
+    msg = read_client.get_wait("test")
+    # time_c = time.time()
+
+    # write_time = int((time_b - time_a) * 1e6)
+    # read_time = int((time_c - time_b) * 1e6)
+    # print(f"write time: {write_time} us,  read time: {read_time} us")
+
+    assert msg == "test_msg"
+
+    stop_flag.set()
+    p_speed_read.join()
+    p_speed_read.kill()
+    read_client.close()
+    write_client.shutdown_server()
+    # time.sleep(_pause)
+    write_client.close()

{dsmq-1.0.0.dist-info → dsmq-1.1.0.dist-info}/METADATA

@@ -1,9 +1,10 @@
 Metadata-Version: 2.4
 Name: dsmq
-Version: 1.0.0
+Version: 1.1.0
 Summary: A dead simple message queue
 License-File: LICENSE
 Requires-Python: >=3.10
+Requires-Dist: pytest>=8.3.4
 Requires-Dist: websockets>=14.1
 Description-Content-Type: text/markdown
 
@@ -86,7 +87,7 @@ p_mq.close()
 1. In the second, run `src/dsmq/example_put_client.py`.
 1. In the third, run `src/dsmq/example_get_client.py`.
 
-Alternatively, if you're on Linux just run `src/dsmq/demo_linux.py`.
+Alternatively, you can run them all at once with `src/dsmq/demo.py`.
 
 ## How it works
 
@@ -150,3 +151,30 @@ connected to the server.
 - returns str, the content of the message. If there was no eligble message
 in the topic, or the topic doesn't yet exist,
 returns `""`.
+
+### `get_wait(topic)`
+
+A variant of `get()` that retries a few times until it gets
+a non-empty message. Adjust internal values `_n_tries` and `_initial_retry`
+to change how persistent it will be.
+
+- `topic` (str)
+- returns str, the content of the message. If there was no eligble message
+in the topic after the allotted number of tries,
+or the topic doesn't yet exist,
+returns `""`.
+
+### `shutdown_server()`
+
+Gracefully shut down the server, through the client connection.
+
+### `close()`
+
+Gracefully shut down the client connection.
+
+# Testing
+
+Run all the tests in `src/dsmq/tests/` with pytest, for example
+```
+uv run pytest
+```

dsmq-1.1.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+dsmq/__init__.py,sha256=YCgbnQAk8YbtHRyMcU0v2O7RdRhPhlT-vS_q40a7Q6g,50
+dsmq/client.py,sha256=2i3BhpQM71j-uZlqiJvzQgywi1y07K1VgQe8i2Koi0g,2070
+dsmq/demo.py,sha256=K53cC5kN7K4kNJlPq7c5OTIMHRCKTo9hYX2aIos57rU,542
+dsmq/example_get_client.py,sha256=PvAsDGEAH1kVBifLVg2rx8ZxnAZmvzVCvZq13VgpLds,301
+dsmq/example_put_client.py,sha256=QxDc3i7KAjjhpwxRRpI0Ke5KTNSPuBf9kkcGyTvUEaw,353
+dsmq/server.py,sha256=iiaa7IoxZy1143vwTO7eMD-kzUKTTVff37E3KPpJ4cU,6010
+dsmq/tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+dsmq/tests/integration_test.py,sha256=TJyx_wsFWeSvYuMF7vCNSGcWstTSLZbuAeKrxrQcj98,5930
+dsmq-1.1.0.dist-info/METADATA,sha256=S72f4YPsNO07F8W4cnrhGGEW-4Eoao9BBh2VP28uziM,4730
+dsmq-1.1.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+dsmq-1.1.0.dist-info/licenses/LICENSE,sha256=3Yu1mAp5VsKmnDtzkiOY7BdmrLeNwwZ3t6iWaLnlL0Y,1071
+dsmq-1.1.0.dist-info/RECORD,,

dsmq-1.0.0.dist-info/RECORD

@@ -1,10 +0,0 @@
-dsmq/__init__.py,sha256=YCgbnQAk8YbtHRyMcU0v2O7RdRhPhlT-vS_q40a7Q6g,50
-dsmq/client.py,sha256=sgkz7iabc8hirq0HRfhqlKCPwHVYgvg2C01aoRbvTSo,768
-dsmq/demo_linux.py,sha256=O5B2DH9pAvCCaYSRE5DV1o2fSx_bXU9hcxTis0ier5I,648
-dsmq/example_get_client.py,sha256=PvAsDGEAH1kVBifLVg2rx8ZxnAZmvzVCvZq13VgpLds,301
-dsmq/example_put_client.py,sha256=QxDc3i7KAjjhpwxRRpI0Ke5KTNSPuBf9kkcGyTvUEaw,353
-dsmq/server.py,sha256=fFM8soN4y8CULX5yk3Temir-pl2K1hp38MOEwqzPN5U,4343
-dsmq-1.0.0.dist-info/METADATA,sha256=TaoyVGkRJasSK-D6UQ9vNTeY07gOH9At7j1rZgiFxp0,4069
-dsmq-1.0.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-dsmq-1.0.0.dist-info/licenses/LICENSE,sha256=3Yu1mAp5VsKmnDtzkiOY7BdmrLeNwwZ3t6iWaLnlL0Y,1071
-dsmq-1.0.0.dist-info/RECORD,,