backend now more robust, threaded writes use locks

cli/src/zshell/subcommands/battery.py

@@ -7,23 +7,32 @@ import logging
 import os
 import re
 import signal
+import struct
 import tempfile
+import threading
 import time
+from collections.abc import Callable
 from dataclasses import asdict, dataclass
 from pathlib import Path
 
 import typer
+from logitech_receiver import hidpp10, hidpp20
+from logitech_receiver.common import Battery, BatteryStatus, Notification
+from logitech_receiver.hidpp10_constants import Registers
+from logitech_receiver.hidpp20_constants import SupportedFeature
 
 app = typer.Typer(
 	help="Read live battery status from Logitech HID++ peripherals."
 )
 
 logger = logging.getLogger("zshell.battery")
+
 DEFAULT_OUTPUT = Path.home() / ".cache" / "zshell" / "battery.json"
+
 _RUNTIME_DIR = Path(
 	os.environ.get("XDG_RUNTIME_DIR") or (Path.home() / ".cache" / "zshell")
 )
-LOCK = _RUNTIME_DIR / "zshell-battery-daemon.lock"
+DEFAULT_LOCK = _RUNTIME_DIR / "zshell-battery-daemon.lock"
 
 
 def _acquire_singleton_lock(path: Path) -> int | None:
@@ -76,6 +85,350 @@ def _find_upower_native_path(serial: str | None) -> str | None:
 	return None
 
 
+def _process_battery_notification(dev, n) -> bool:
+
+	if not getattr(dev, "isDevice", False):
+		return False
+
+	if int(n.sub_id) & 0x80:
+		return False
+	if n.sub_id == Notification.NO_OPERATION:
+		return False
+
+	if getattr(dev, "protocol", None) is not None and dev.protocol < 2.0:
+		if n.sub_id in (Registers.BATTERY_STATUS, Registers.BATTERY_CHARGE):
+			if n.data[-1:] != b"\x00":
+				return False
+			data = bytes([n.address]) + n.data
+			dev.set_battery_info(hidpp10.parse_battery_status(n.sub_id, data))
+			return True
+		return False
+
+	features = getattr(dev, "features", None)
+	if not features:
+		return False
+
+	try:
+		feature = features.get_feature(n.sub_id)
+	except Exception:
+		return False
+
+	if feature == SupportedFeature.BATTERY_STATUS:
+		if n.address == 0x00:
+			dev.set_battery_info(hidpp20.decipher_battery_status(n.data)[1])
+			return True
+		return False
+
+	if feature == SupportedFeature.BATTERY_VOLTAGE:
+		if n.address == 0x00:
+			dev.set_battery_info(hidpp20.decipher_battery_voltage(n.data)[1])
+			return True
+		return False
+
+	if feature == SupportedFeature.UNIFIED_BATTERY:
+		if n.address == 0x00:
+			dev.set_battery_info(hidpp20.decipher_battery_unified(n.data)[1])
+			return True
+		return False
+
+	if feature == SupportedFeature.ADC_MEASUREMENT:
+		if n.address == 0x00:
+			result = hidpp20.decipher_adc_measurement(n.data)
+			if result:
+				dev.set_battery_info(result[1])
+				return True
+		return False
+
+	if feature == SupportedFeature.CENTURION_BATTERY_SOC:
+		dev.set_battery_info(hidpp20.decipher_battery_centurion(n.data)[1])
+		return True
+
+	if feature == SupportedFeature.SOLAR_DASHBOARD:
+		if n.data[5:9] == b"GOOD":
+			charge, lux, adc = struct.unpack("!BHH", n.data[:5])
+			status = BatteryStatus.DISCHARGING
+			if n.address == 0x10 and lux > 200:
+				status = BatteryStatus.RECHARGING
+			dev.set_battery_info(Battery(charge, None, status, None, lux))
+			return True
+		return False
+
+	return False
+
+
+def _import_listener_deps():
+	try:
+		from logitech_receiver import base, device, receiver
+		from logitech_receiver.listener import EventsListener
+	except ImportError as e:
+		raise RuntimeError(
+			"logitech_receiver isn't importable -- install it with `pip install solaar`."
+		) from e
+	except ValueError as e:
+		raise RuntimeError(
+			"logitech_receiver failed to import because a GTK3 typelib is missing "
+			f"({e}). Install the same GTK3 + PyGObject packages Solaar's GUI needs "
+			"(e.g. python3-gi + gir1.2-gtk-3.0 on Debian/Ubuntu, or the gtk3/"
+			"python3-gobject equivalents on your distro)."
+		) from e
+
+	class _BatteryEventsListener(EventsListener):
+		def __init__(self, receiver_or_device, on_change):
+			super().__init__(receiver_or_device, self._handle_notification)
+			self._on_change = on_change
+
+		def has_started(self):
+
+			if not self.receiver.isDevice:
+				try:
+					self.receiver.notification_flags = (
+						self.receiver.enable_connection_notifications()
+					)
+					self.receiver.notify_devices()
+				except Exception:
+					logger.exception(
+						"failed enabling notifications for %s", self.receiver
+					)
+
+		def _handle_notification(self, n):
+			if self.receiver.isDevice:
+				try:
+					_process_battery_notification(self.receiver, n)
+				except Exception:
+					logger.exception("processing %s for %s", n, self.receiver)
+				self._on_change(self.receiver)
+				return
+
+			if n.devnumber == 0xFF:
+				try:
+					_process_battery_notification(self.receiver, n)
+				except Exception:
+					logger.exception("processing receiver notification %s", n)
+				self._on_change(self.receiver)
+				return
+
+			if not (0 < n.devnumber <= 16):
+				logger.warning(
+					"unexpected device number %s in %s", n.devnumber, n
+				)
+				return
+
+			try:
+				dev = self.receiver[n.devnumber]
+			except Exception:
+				logger.exception(
+					"resolving device number %s on %s",
+					n.devnumber,
+					self.receiver,
+				)
+				return
+
+			if not dev:
+				logger.warning(
+					"%s: received %s for invalid device %d",
+					self.receiver,
+					n,
+					n.devnumber,
+				)
+				return
+
+			try:
+				_process_battery_notification(dev, n)
+			except Exception:
+				logger.exception("processing %s for %s", n, dev)
+
+			self._on_change(dev)
+
+	return base, device, receiver, _BatteryEventsListener
+
+
+def _start_hotplug_watcher(on_hotplug: Callable[[], None]):
+	try:
+		import pyudev
+	except ImportError:
+		logger.warning(
+			"pyudev isn't importable -- hotplugged devices will only be "
+			"noticed on the next --rescan-interval tick, not instantly. "
+			"Install it with `pip install pyudev` (it's a core Solaar "
+			"dependency, so it's normally already present)."
+		)
+		return None
+
+	debounce_lock = threading.Lock()
+	debounce_state = {"timer": None}
+	DEBOUNCE_SECONDS = 0.5
+
+	def _fire():
+		with debounce_lock:
+			debounce_state["timer"] = None
+		on_hotplug()
+
+	def _handle_event(_device):
+		if _device.action not in ("add", "remove"):
+			return
+		with debounce_lock:
+			existing = debounce_state["timer"]
+			if existing is not None:
+				existing.cancel()
+			timer = threading.Timer(DEBOUNCE_SECONDS, _fire)
+			timer.daemon = True
+			debounce_state["timer"] = timer
+			timer.start()
+
+	context = pyudev.Context()
+	monitor = pyudev.Monitor.from_netlink(context)
+	monitor.filter_by(subsystem="hidraw")
+	observer = pyudev.MonitorObserver(monitor, callback=_handle_event)
+	observer.start()
+	return observer
+
+
+class _ListenerRegistry:
+	def __init__(self, on_change):
+		self._on_change = on_change
+		self._entries = {}
+		self._last_change: dict[int, float] = {}
+		self._lock = threading.Lock()
+		self._rescan_lock = threading.Lock()
+		self._min_rescan_interval = 1.0
+		self._last_rescan_time = 0.0
+		self._hidpp_io_lock = threading.RLock()
+
+	def _touch_and_notify(self, dev_or_receiver, *args, **kwargs):
+		with self._lock:
+			self._last_change[id(dev_or_receiver)] = time.time()
+		self._on_change(dev_or_receiver, *args, **kwargs)
+
+	def last_change(self, dev) -> float:
+		with self._lock:
+			return self._last_change.get(id(dev), 0.0)
+
+	def rescan(self):
+		now = time.monotonic()
+		if now - self._last_rescan_time < self._min_rescan_interval:
+			return
+		if not self._rescan_lock.acquire(blocking=False):
+			return
+		try:
+			self._last_rescan_time = now
+			with self._hidpp_io_lock:
+				base, device, receiver, ListenerClass = _import_listener_deps()
+
+				with self._lock:
+					dead = [
+						p
+						for p, (_obj, listener) in self._entries.items()
+						if not listener._active
+					]
+					for p in dead:
+						del self._entries[p]
+					known_paths = set(self._entries.keys())
+
+				for dev_info in base.receivers_and_devices():
+					if dev_info.path in known_paths:
+						continue
+					try:
+						if dev_info.isDevice:
+							obj = device.create_device(base, dev_info)
+						else:
+							obj = receiver.create_receiver(base, dev_info)
+					except OSError as e:
+						if e.errno == 13:
+							logger.error(
+								"permission denied opening %s -- check the Solaar udev rule "
+								"(rules.d/42-logitech-unify-permissions.rules) is installed",
+								dev_info.path,
+							)
+						else:
+							logger.exception("failed opening %s", dev_info)
+						continue
+					except Exception:
+						logger.exception("failed opening %s", dev_info)
+						continue
+
+					if obj is None:
+						continue
+
+					listener = ListenerClass(obj, self._touch_and_notify)
+					listener.start()
+
+					with self._lock:
+						self._entries[dev_info.path] = (obj, listener)
+						self._last_change[id(obj)] = time.time()
+
+					logger.info(
+						"listening on %s (%s)",
+						dev_info.path,
+						"device" if dev_info.isDevice else "receiver",
+					)
+		finally:
+			self._rescan_lock.release()
+
+	def known_devices(self):
+		with self._lock:
+			entries = list(self._entries.values())
+
+		result = []
+		for obj, _listener in entries:
+			if obj.isDevice:
+				result.append(obj)
+			else:
+				result.extend(list(obj))
+		return result
+
+	def stop(self):
+		with self._lock:
+			entries = list(self._entries.values())
+		for _obj, listener in entries:
+			try:
+				listener.stop()
+			except Exception:
+				logger.exception("error stopping listener")
+		for _obj, listener in entries:
+			listener.join(timeout=2.0)
+
+
+def _dedupe_by_identity(
+	pairs: list[tuple[object, DeviceBattery]],
+	freshness: Callable[[object], float],
+) -> list[DeviceBattery]:
+	passthrough: list[DeviceBattery] = []
+	best: dict[str, tuple[float, DeviceBattery]] = {}
+
+	for dev, entry in pairs:
+		key = entry.serial
+		if not key:
+			passthrough.append(entry)
+			continue
+		ts = freshness(dev)
+		current = best.get(key)
+		if current is None or ts >= current[0]:
+			best[key] = (ts, entry)
+
+	return passthrough + [entry for _ts, entry in best.values()]
+
+
+def _snapshot(registry: _ListenerRegistry) -> list[DeviceBattery]:
+	pairs: list[tuple[object, DeviceBattery]] = []
+
+	with registry._hidpp_io_lock:
+		for dev in registry.known_devices():
+			info = dev.battery_info
+			if info is None:
+				try:
+					if dev.ping():
+						info = dev.battery()
+				except Exception:
+					logger.debug(
+						"initial battery() failed for %s", dev, exc_info=True
+					)
+					info = None
+			entry = _to_device_battery(dev, info)
+			if entry is not None:
+				pairs.append((dev, entry))
+	return _dedupe_by_identity(pairs, registry.last_change)
+
+
 def _clean_json(path: Path) -> None:
 	_write_json_atomic(path, [])
 
@@ -88,7 +441,6 @@ def _iter_open_devices():
 			"logitech_receiver isn't importable -- install it with `pip install solaar`."
 		) from e
 	except ValueError as e:
-		# gi.require_version() raises ValueError if the GTK3 typelib isn't installed.
 		raise RuntimeError(
 			"logitech_receiver failed to import because a GTK3 typelib is missing "
 			f"({e}). Install the same GTK3 + PyGObject packages Solaar's GUI needs "
@@ -103,7 +455,7 @@ def _iter_open_devices():
 			else:
 				d = receiver.create_receiver(base, dev_info)
 		except OSError as e:
-			if e.errno == 13:  # EACCES
+			if e.errno == 13:
 				logger.error(
 					"permission denied opening %s -- check the Solaar udev rule "
 					"(rules.d/42-logitech-unify-permissions.rules) is installed "
@@ -123,7 +475,31 @@ def _iter_open_devices():
 		if d.isDevice:
 			yield d
 		else:
-			yield from d  # a receiver: walk its currently paired devices
+			yield from d
+
+
+def _to_device_battery(dev, battery) -> DeviceBattery | None:
+	if battery is None:
+		return None
+
+	percentage = battery.level / 100 if isinstance(battery.level, int) else None
+	state = battery.status.name.lower() if battery.status is not None else None
+
+	serial = (
+		getattr(dev, "serial", None) or getattr(dev, "unitId", None) or None
+	)
+
+	return DeviceBattery(
+		model=dev.name or dev.codename or "Unknown device",
+		nativePath=_find_upower_native_path(serial),
+		serial=serial,
+		type=str(dev.kind) if dev.kind is not None else None,
+		wired=dev.receiver is None,
+		percentage=percentage,
+		state=state,
+		charging=battery.charging(),
+		isPresent=battery.ok(),
+	)
 
 
 def _read_battery(dev) -> DeviceBattery | None:
@@ -140,37 +516,23 @@ def _read_battery(dev) -> DeviceBattery | None:
 		logger.debug("battery() failed for %s", dev, exc_info=True)
 		return None
 
-	if battery is None:
-		return None
-
-	percentage = battery.level / 100 if isinstance(battery.level, int) else None
-	state = battery.status.name.lower() if battery.status is not None else None
-	serial = getattr(dev, "serial", None)
-
-	return DeviceBattery(
-		model=dev.name or dev.codename or "Unknown device",
-		nativePath=_find_upower_native_path(serial),
-		serial=serial,
-		type=str(dev.kind) if dev.kind is not None else None,
-		wired=dev.receiver is None,
-		percentage=percentage,
-		state=state,
-		charging=battery.charging(),
-		isPresent=battery.ok(),
-	)
+	return _to_device_battery(dev, battery)
 
 
 def poll_once() -> list[DeviceBattery]:
-	results = []
+	pairs: list[tuple[object, DeviceBattery]] = []
 	for dev in _iter_open_devices():
 		try:
 			info = _read_battery(dev)
 			if info is not None:
-				results.append(info)
+				pairs.append((dev, info))
 		finally:
 			with contextlib.suppress(Exception):
 				dev.close()
-	return results
+
+	return _dedupe_by_identity(
+		pairs, lambda dev: 1.0 if dev.receiver is None else 0.0
+	)
 
 
 def _write_json_atomic(path: Path, results: list[DeviceBattery]) -> None:
@@ -185,7 +547,7 @@ def _write_json_atomic(path: Path, results: list[DeviceBattery]) -> None:
 	try:
 		with os.fdopen(fd, "w") as f:
 			json.dump(payload, f, indent=2)
-		os.replace(tmp_path, path)  # atomic rename on the same filesystem
+		os.replace(tmp_path, path)
 	except Exception:
 		Path(tmp_path).unlink(missing_ok=True)
 		raise
@@ -193,8 +555,18 @@ def _write_json_atomic(path: Path, results: list[DeviceBattery]) -> None:
 
 @app.command()
 def daemon(
-	interval: float = typer.Option(
-		15.0, "--interval", "-i", help="Seconds between polls."
+	rescan_interval: float = typer.Option(
+		60.0,
+		"--rescan-interval",
+		"-i",
+		help=(
+			"Seconds between fallback rescans. This is NOT a battery poll "
+			"interval -- battery updates are event-driven and written as "
+			"soon as a device reports a change. New/removed devices "
+			"(e.g. plugging a keyboard in to charge) are normally noticed "
+			"within milliseconds via a udev hotplug watcher; this interval "
+			"is only a safety net in case a udev event is ever missed."
+		),
 	),
 	out: Path = typer.Option(
 		DEFAULT_OUTPUT,
@@ -202,22 +574,86 @@ def daemon(
 		"-o",
 		help="Where to write the JSON status file.",
 	),
-	lock: Path = typer.Option(
-		LOCK, "--lock-file", help="Path to daemon lock file"
+	lock_file: Path = typer.Option(
+		DEFAULT_LOCK,
+		"--lock-file",
+		help="Path used to ensure only one daemon runs at a time.",
+	),
+	verbose: bool = typer.Option(
+		False,
+		"--verbose",
+		"-v",
+		help="Log every listener start and battery-change write to stderr.",
 	),
 ):
-	lock_fd = _acquire_singleton_lock(lock)
+	lock_fd = _acquire_singleton_lock(lock_file)
 	if lock_fd is None:
 		try:
-			holder_pid = lock.read_text().strip()
+			holder_pid = lock_file.read_text().strip()
 		except OSError:
 			holder_pid = "unknown"
 		typer.echo(
-			f"error: a battery daemon is already running (pid: {holder_pid}, lock: {lock})",
+			f"error: a battery daemon is already running (pid {holder_pid}, lock: {lock_file})",
 			err=True,
 		)
 		raise typer.Exit(code=1)
 
+	if verbose:
+		logging.basicConfig(
+			level=logging.INFO, format="%(asctime)s %(name)s: %(message)s"
+		)
+
+	write_lock = threading.Lock()
+	last_payload = None
+
+	pending_write_timer: list[threading.Timer | None] = [None]
+	pending_write_lock = threading.Lock()
+	WRITE_SETTLE_SECONDS = 0.3
+
+	def write_if_changed():
+		nonlocal last_payload
+		with write_lock:
+			results = _snapshot(registry)
+			payload = [asdict(r) for r in results]
+			if payload != last_payload:
+				_write_json_atomic(out, results)
+				last_payload = payload
+				logger.info(
+					"battery status changed, wrote %d device(s)", len(results)
+				)
+
+	def write_if_changed_coalesced():
+		with pending_write_lock:
+			existing = pending_write_timer[0]
+			if existing is not None:
+				existing.cancel()
+			timer = threading.Timer(WRITE_SETTLE_SECONDS, write_if_changed)
+			timer.daemon = True
+			pending_write_timer[0] = timer
+			timer.start()
+
+	def on_change(_device_or_receiver, alert=None, reason=None):
+		write_if_changed_coalesced()
+
+	registry = _ListenerRegistry(on_change)
+
+	def rescan_and_write():
+		try:
+			registry.rescan()
+		except Exception:
+			logger.exception("rescan failed")
+		write_if_changed_coalesced()
+
+	try:
+		registry.rescan()
+	except RuntimeError as e:
+		typer.echo(f"error: {e}", err=True)
+		raise typer.Exit(code=1) from None
+
+	write_if_changed()
+
+	hotplug_observer = _start_hotplug_watcher(rescan_and_write)
+
 	running = True
 
 	def _stop(signum, frame):
@@ -227,32 +663,27 @@ def daemon(
 	signal.signal(signal.SIGINT, _stop)
 	signal.signal(signal.SIGTERM, _stop)
 
-	typer.echo(f"Polling every {interval}s, writing to {out}")
+	typer.echo(
+		f"Listening for battery events, writing to {out} (lock: {lock_file})"
+	)
 
-	last_payload = None
 	try:
 		while running:
-			try:
-				results = poll_once()
-				payload = [asdict(r) for r in results]
-				if payload != last_payload:
-					_write_json_atomic(out, results)
-					last_payload = payload
-			except RuntimeError as e:
-				typer.echo(f"error: {e}", err=True)
-				raise typer.Exit(code=1) from None
-			except Exception:
-				logger.exception("poll cycle failed")
-
-			remaining = interval
+			remaining = rescan_interval
 			while running and remaining > 0:
 				step = min(0.2, remaining)
 				time.sleep(step)
 				remaining -= step
+			if running:
+				rescan_and_write()
 	finally:
+		if hotplug_observer is not None:
+			with contextlib.suppress(Exception):
+				hotplug_observer.stop()
 		with contextlib.suppress(Exception):
 			_clean_json(out)
 
+	registry.stop()
 	typer.echo("Stopped.")