From b4ed4f8e6d284578a3720131735301b0f85cbdcf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Przemys=C5=82aw=20Barcicki?= <pbarcicki@antmicro.com>
Date: Wed, 28 Jan 2026 11:05:16 +0100
Subject: [PATCH] [#85496] Split socket function

---
 sargraph.py | 11 +++++------
 watch.py    | 24 ++++++++++++------------
 2 files changed, 17 insertions(+), 18 deletions(-)

diff --git a/sargraph.py b/sargraph.py
index ceafd16..cec6e1c 100755
--- a/sargraph.py
+++ b/sargraph.py
@@ -7,7 +7,6 @@
 
 import argparse
 import sys
-import time
 
 import graph
 import watch
@@ -31,10 +30,11 @@ parser.add_argument('-p',      action='store_true',
 args = parser.parse_args()
 
 def send(session: str, message: str):
-    sock, socket_path = watch.get_socket(session)
+    socket_path = watch.get_socket_path(session)
     if not file_exists(socket_path):
         fail(f"Session '{session}' does not exist")
 
+    sock = watch.get_bound_socket(socket_path)
     sock.connect(socket_path)
     sock.send(message.encode("utf-8"))
     sock.close()
@@ -67,8 +67,7 @@ if args.name != "data":
 
 # Check if a command was provided, if that session exists, yell at user for lack of commands, else spawn
 if len(args.command) == 0:
-    sock, socket_path = watch.get_socket(args.session)
-    if file_exists(socket_path):
+    if file_exists(watch.get_socket_path(args.session)):
         fail("Command not provided")
     
     else:
@@ -76,7 +75,7 @@ if len(args.command) == 0:
         create_session()
 
 if args.command[0] == "start":
-    sock, socket_path = watch.get_socket(args.session)
+    socket_path = watch.get_socket_path(args.session)
     if file_exists(socket_path):
         fail("Session with this name already exists")
     
@@ -95,7 +94,7 @@ if args.command[0] == "start":
     fail("Session did not start")
 
 elif args.command[0] == "stop":
-    _, socket_path = watch.get_socket(args.session)
+    socket_path = watch.get_socket_path(args.session)
     
     print(f"Terminating sargraph session '{args.session}'")
     if len(args.command) < 2:
diff --git a/watch.py b/watch.py
index 008a870..9529faf 100644
--- a/watch.py
+++ b/watch.py
@@ -101,16 +101,16 @@ def read_iface_stats(iface):
         tx = scan(r"(\d+)", int, f.readline())
     return rx, tx
 
-def get_socket(session):
-    if is_windows():
-        path = fr"\\.\pipe\sargraph-{session}"
-    else:
-        path = f"/tmp/sargraph-{session}.sock"
-    sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
-    return sock, path
+def get_socket_path(session):
+    return fr"\\.\pipe\sargraph-{session}" if is_windows() else f"/tmp/sargraph-{session}.sock"
 
+def get_bound_socket(sock_path):
+    sock = socket.socket(socket.AF_UNIX, socket.SOCK_DGRAM)
+    sock.bind(sock_path)
+    return sock
 
 class Watcher(abc.ABC):
+    sock: socket.socket
 
     def __init__(self, session, fsdev, iface, tmpfs_color, other_cache_color, udp=None, udp_cookie=None):
         super().__init__()
@@ -129,7 +129,7 @@ class Watcher(abc.ABC):
         file_handler.setFormatter(logging.Formatter("%(message)s"))
         self.logger.addHandler(file_handler)
 
-        self.socket, self.socket_path = get_socket(self.session)
+        self.socket_path = get_socket_path(session)
 
         # Was a graph already produced by save command from sargraph?
         self.dont_plot = False
@@ -160,7 +160,7 @@ class Watcher(abc.ABC):
         self.die = True
 
     def recv_data(self) -> str:
-        data = self.socket.recv(1 << 10)  # 1024 bytes should be enough
+        data = self.sock.recv(1 << 10)  # 1024 bytes should be enough
         return data.decode("utf-8").replace("\n", "").strip()
 
     # Add a summary comment to 'data.txt'
@@ -237,7 +237,7 @@ class Watcher(abc.ABC):
         self.logger.info(msg)
 
     def start(self):
-        self.socket.bind(self.socket_path)
+        self.sock = get_bound_socket(self.socket_path)
         signal.signal(signal.SIGTERM, self.kill_handler)
 
         try:
@@ -387,7 +387,7 @@ class SarWatcher(Watcher):
         curr_gpu_util = 0
         curr_gpu_mem = 0
 
-        socket_fd = self.socket.fileno()
+        socket_fd = self.sock.fileno()
 
         while 1:
             # Await sar output or a command sent from command handler in sargraph.py
@@ -621,7 +621,7 @@ class PsUtilWatcher(Watcher):
         thread.start()
 
         self.initialize()
-        socket_fd = self.socket.fileno()
+        socket_fd = self.sock.fileno()
 
         while 1:
             # Await sar output or a command sent from command handler in sargraph.py