Annotate types

Author: ben-grande

qubes/qmemman/algo.py

@@ -19,6 +19,7 @@
 #
 
 import logging
+from typing import Optional
 
 # These defaults can be overridden by QMemmanServer with values from config
 # file.
@@ -33,7 +34,7 @@
 log = logging.getLogger("qmemman.daemon.algo")
 
 
-def sanitize_and_parse_meminfo(untrusted_meminfo):
+def sanitize_and_parse_meminfo(untrusted_meminfo) -> Optional[int]:
     # Untrusted meminfo size is read from xenstore, thus its size is limited
     # and splits do not require excessive memory.
     if not untrusted_meminfo:
@@ -43,14 +44,14 @@ def sanitize_and_parse_meminfo(untrusted_meminfo):
     return int(untrusted_meminfo) * 1024
 
 
-def refresh_meminfo_for_domain(domain, untrusted_xenstore_key):
+def refresh_meminfo_for_domain(domain, untrusted_xenstore_key) -> None:
     """
     Called when a domain updates its 'meminfo' xenstore key.
     """
     domain.mem_used = sanitize_and_parse_meminfo(untrusted_xenstore_key)
 
 
-def pref_mem(domain):
+def pref_mem(domain) -> int:
     # As dom0 must have large cache for vbds, give it a special boost.
     mem_used = domain.mem_used * CACHE_FACTOR
     if domain.domid == "0":
@@ -59,7 +60,7 @@ def pref_mem(domain):
     return int(max(min(mem_used, domain.mem_max), MIN_PREFMEM))
 
 
-def needed_mem(domain):
+def needed_mem(domain) -> int:
     # Do not change. In balance(), "distribute total_available_mem
     # proportionally to pref_mem" relies on this exact formula.
     ret = pref_mem(domain) - domain.mem_actual
@@ -69,7 +70,7 @@ def needed_mem(domain):
 # Prepare list of (domain, mem_target) pairs that need to be passed to "xm
 # memset" equivalent in order to obtain "mem_size".
 # Returns empty list when the request cannot be satisfied.
-def balloon(mem_size, dom_dict):
+def balloon(mem_size, dom_dict) -> list:
     log.debug(
         "balloon(mem_size={!r}, dom_dict={!r})".format(mem_size, dom_dict)
     )
@@ -212,7 +213,7 @@ def balance_when_low_on_mem(
 # Get memory information.
 # Called before and after domain balances.
 # Return a dictionary of various memory data points.
-def mem_info(xen_free_mem, dom_dict):
+def mem_info(xen_free_mem, dom_dict) -> dict:
     log.debug(
         "mem_info(xen_free_mem={!r}, dom_dict={!r})".format(
             xen_free_mem, dom_dict
@@ -264,7 +265,7 @@ def mem_info(xen_free_mem, dom_dict):
 # Called when one of domains update its 'meminfo' xenstore key.
 # Return the list of (domain, mem_target) pairs to be passed to "xm memset"
 # equivalent
-def balance(xen_free_mem, dom_dict):
+def balance(xen_free_mem, dom_dict) -> dict:
     log.debug(
         "balance(xen_free_mem={!r}, dom_dict={!r})".format(
             xen_free_mem, dom_dict

qubes/qmemman/client.py

@@ -19,13 +19,14 @@
 
 import socket
 import fcntl
+from typing import Optional
 
 
 class QMemmanClient:
-    def __init__(self):
-        self.sock = None
+    def __init__(self) -> None:
+        self.sock: Optional[socket.socket] = None
 
-    def request_mem(self, amount):
+    def request_mem(self, amount) -> bool:
         self.sock = socket.socket(socket.AF_UNIX)
         flags = fcntl.fcntl(self.sock.fileno(), fcntl.F_GETFD)
         flags |= fcntl.FD_CLOEXEC
@@ -35,5 +36,6 @@ def request_mem(self, amount):
         received = self.sock.recv(1024).strip()
         return bool(received == b"OK")
 
-    def close(self):
+    def close(self) -> None:
+        assert isinstance(self.sock, socket.socket)
         self.sock.close()

qubes/qmemman/domainstate.py

@@ -18,20 +18,29 @@
 # You should have received a copy of the GNU General Public
 # License along with this library; if not, see <https://www.gnu.org/licenses/>.
 
+from typing import Optional
+
 
 class DomainState:  # pylint: disable=too-few-public-methods
-    def __init__(self, domid):
-        self.mem_current = 0  # the current memory size
-        self.mem_actual = None  # the current memory allocation (what VM
-        # is using or can use at any time)
-        self.mem_max = None  # the maximum memory size
-        self.mem_used = None  # used memory, computed based on meminfo
-        self.domid = domid  # domain id
-        self.last_target = 0  # the last memset target
-        self.use_hotplug = False  # use memory hotplug for mem-set
-        self.no_progress = False  # no react to memset
-        self.slow_memset_react = False  # slow react to memset (after few
-        # tries still above target)
+    def __init__(self, domid) -> None:
+        # Current memory size.
+        self.mem_current: int = 0
+        # Current memory allocation (what VM is using or can use at any time).
+        self.mem_actual: Optional[int] = None
+        # Maximum memory size.
+        self.mem_max: Optional[int] = None
+        # Used memory, computed based on meminfo.
+        self.mem_used: Optional[int] = None
+        # Domain ID.
+        self.domid: str = domid
+        # Last memset target.
+        self.last_target: int = 0
+        # Use memory hotplug for mem-set.
+        self.use_hotplug: bool = False
+        # No reaction to memset.
+        self.no_progress: bool = False
+        # Slow react to memset (after few tries still above target).
+        self.slow_memset_react: bool = False
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return self.__dict__.__repr__()

qubes/qmemman/systemstate.py

@@ -24,6 +24,7 @@
 import time
 import xen.lowlevel  # pylint: disable=import-error
 from pathlib import Path
+from typing import Optional
 
 import qubes.qmemman
 from qubes.qmemman.domainstate import DomainState
@@ -39,16 +40,16 @@
 
 
 class SystemState:
-    def __init__(self):
+    def __init__(self) -> None:
         self.log = logging.getLogger("qmemman.systemstate")
         self.log.debug("SystemState()")
 
-        self.dom_dict = {}
-        self.xc = None
-        self.xs = None
-        self.all_phys_mem = 0
+        self.dom_dict: dict[str, DomainState] = {}
+        self.xc: xen.lowlevel.xc.xc = None
+        self.xs: xen.lowlevel.xs.xs = None
+        self.all_phys_mem: int = 0
 
-    def init(self):
+    def init(self) -> None:
         self.xc = xen.lowlevel.xc.xc()
         self.xs = xen.lowlevel.xs.xs()
         # Overhead of per-page Xen structures, taken from OpenStack
@@ -63,28 +64,28 @@ def init(self):
         except xen.lowlevel.xc.Error:
             pass
 
-    def xs_wrapper(self, key, value=None):
+    def xs_wrapper(self, key, value=None) -> Optional[str]:
         if value:
             self.xs.write("", key, value)
             return None
         return self.xs.read("", key)
 
-    def get_xs_path(self, domid, key):
+    def get_xs_path(self, domid, key) -> str:
         return "/local/domain/" + str(domid) + "/memory/" + key
 
-    def add_domain(self, domid):
+    def add_domain(self, domid) -> None:
         self.log.debug("add_domain(domid={!r})".format(domid))
         self.dom_dict[domid] = DomainState(domid)
         # TODO: move to DomainState.__init__
         target_str = self.xs_wrapper(self.get_xs_path(domid, "target"))
         if target_str:
             self.dom_dict[domid].last_target = int(target_str) * 1024
 
-    def del_domain(self, domid):
+    def del_domain(self, domid) -> None:
         self.log.debug("del_domain(domid={!r})".format(domid))
         self.dom_dict.pop(domid)
 
-    def get_free_xen_mem(self):
+    def get_free_xen_mem(self) -> int:
         xen_free = int(
             self.xc.physinfo()["free_memory"] * 1024 * MEM_OVERHEAD_FACTOR
         )
@@ -116,7 +117,9 @@ def get_free_xen_mem(self):
         return xen_free - assigned_but_unused
 
     # Refresh information on memory assigned to all domains
-    def refresh_mem_actual(self):
+    def refresh_mem_actual(self) -> None:
+        if self.xc is None:
+            raise AttributeError("self.xc has no methods")
         for domain in self.xc.domain_getinfo():
             domid = str(domain["domid"])
             if domid in self.dom_dict:
@@ -149,7 +152,7 @@ def refresh_mem_actual(self):
                     # the only possible case of nonexisting memory/static-max
                     # is dom0, see #307
 
-    def clear_outdated_error_markers(self):
+    def clear_outdated_error_markers(self) -> None:
         # Clear outdated errors.
         for _, dom in self.dom_dict.items():
             if dom.mem_used is None:
@@ -159,6 +162,7 @@ def clear_outdated_error_markers(self):
             #  - VM request more memory than it has assigned
             # The second condition avoids starving a VM, even when there is
             # some free memory available.
+            assert isinstance(dom.mem_actual, int)
             if (
                 dom.mem_actual <= dom.last_target + XEN_FREE_MEM_LEFT / 2
                 or dom.mem_actual < qubes.qmemman.algo.pref_mem(dom)
@@ -168,7 +172,7 @@ def clear_outdated_error_markers(self):
 
     # The below works (and is fast), but then 'xm list' shows unchanged memory
     # value.
-    def mem_set(self, domid, val):
+    def mem_set(self, domid, val) -> None:
         self.log.info("mem-set domain {} to {}".format(domid, val))
         dom = self.dom_dict[domid]
         dom.last_target = val
@@ -196,7 +200,7 @@ def mem_set(self, domid, val):
 
     # This is called at the end of ballooning, when we have Xen free mem
     # already, make sure that past mem_set will not decrease Xen free mem.
-    def inhibit_balloon_up(self):
+    def inhibit_balloon_up(self) -> None:
         self.log.debug("inhibit_balloon_up()")
         for domid, dom in self.dom_dict.items():
             if (
@@ -210,10 +214,10 @@ def inhibit_balloon_up(self):
 
     # Perform memory ballooning, across all domains, to add "mem_size" to Xen
     # free memory
-    def do_balloon(self, mem_size):
+    def do_balloon(self, mem_size) -> bool:
         self.log.info("do_balloon(mem_size={!r})".format(mem_size))
         niter = 0
-        prev_mem_actual = None
+        prev_mem_actual: dict[str, Optional[int]] = {}
 
         for _, dom in self.dom_dict.items():
             dom.no_progress = False
@@ -244,33 +248,30 @@ def do_balloon(self, mem_size):
             ):
                 return False
             xenfree_ring[ring_slot] = xenfree
-            if prev_mem_actual is not None:
-                for domid, dom in prev_mem_actual.items():
-                    if prev_mem_actual[domid] == dom.mem_actual:
-                        # domain not responding to memset requests, remove it
-                        #  from donors
-                        dom.no_progress = True
-                        self.log.info(
-                            "domain {} stuck at {}".format(
-                                domid, dom.mem_actual
-                            )
-                        )
+            for domid, prev_mem in prev_mem_actual.items():
+                dom = self.dom_dict[domid]
+                if prev_mem == dom.mem_actual:
+                    # domain not responding to memset requests, remove it
+                    #  from donors
+                    dom.no_progress = True
+                    self.log.info(
+                        "domain {} stuck at {}".format(domid, dom.mem_actual)
+                    )
             memset_reqs = qubes.qmemman.algo.balloon(
                 mem_size + XEN_FREE_MEM_LEFT - xenfree, self.dom_dict
             )
             self.log.info("memset_reqs={!r}".format(memset_reqs))
             if len(memset_reqs) == 0:
                 return False
             prev_mem_actual = {}
-            for req in memset_reqs:
-                dom, mem = req
-                self.mem_set(dom, mem)
-                prev_mem_actual[dom] = self.dom_dict[dom].mem_actual
+            for domid, memset in memset_reqs:
+                self.mem_set(domid, memset)
+                prev_mem_actual[domid] = self.dom_dict[domid].mem_actual
             self.log.debug("sleeping for {} s".format(BALOON_DELAY))
             time.sleep(BALOON_DELAY)
             niter = niter + 1
 
-    def refresh_meminfo(self, domid, untrusted_meminfo_key):
+    def refresh_meminfo(self, domid, untrusted_meminfo_key) -> None:
         self.log.debug(
             "refresh_meminfo(domid={}, untrusted_meminfo_key={!r})".format(
                 domid, untrusted_meminfo_key
@@ -284,7 +285,7 @@ def refresh_meminfo(self, domid, untrusted_meminfo_key):
 
     # Is the computed balance request big enough so that we do not trash with
     # small adjustments.
-    def is_balance_req_significant(self, memset_reqs, xenfree):
+    def is_balance_req_significant(self, memset_reqs, xenfree) -> bool:
         self.log.debug(
             "is_balance_req_significant(memset_reqs={}, xenfree={})".format(
                 memset_reqs, xenfree
@@ -298,19 +299,18 @@ def is_balance_req_significant(self, memset_reqs, xenfree):
             self.log.debug("xenfree is too low, returning")
             return True
 
-        for req in memset_reqs:
-            dom, mem = req
-            last_target = self.dom_dict[dom].last_target
-            mem_change = mem - last_target
+        for domid, memset in memset_reqs:
+            last_target = self.dom_dict[domid].last_target
+            mem_change = memset - last_target
             total_mem_transfer += abs(mem_change)
-            pref = qubes.qmemman.algo.pref_mem(self.dom_dict[dom])
+            pref = qubes.qmemman.algo.pref_mem(self.dom_dict[domid])
 
             if (
                 0 < last_target < pref
                 and mem_change > MIN_MEM_CHANGE_WHEN_UNDER_PREF
             ):
                 self.log.info(
-                    "dom {} is below pref, allowing balance".format(dom)
+                    "dom {} is below pref, allowing balance".format(domid)
                 )
                 return True
 
@@ -321,7 +321,7 @@ def is_balance_req_significant(self, memset_reqs, xenfree):
         self.log.debug("is_balance_req_significant return {}".format(ret))
         return ret
 
-    def print_stats(self, xenfree, memset_reqs):
+    def print_stats(self, xenfree, memset_reqs) -> None:
         for domid, dom in self.dom_dict.items():
             if dom.mem_used is not None:
                 self.log.info(
@@ -340,14 +340,17 @@ def print_stats(self, xenfree, memset_reqs):
             "stat: xenfree={} memset_reqs={}".format(xenfree, memset_reqs)
         )
 
-    def debug_stuck_balance(self, stuck_domid, memset_reqs, prev_mem_actual):
+    def debug_stuck_balance(
+        self, stuck_domid, memset_reqs, prev_mem_actual
+    ) -> None:
         for req in memset_reqs:
             domid, mem = req
             if domid == stuck_domid:
                 # All donors have been processed.
                 break
             dom = self.dom_dict[domid]
             # Allow some small margin.
+            assert isinstance(dom.mem_actual, int)
             if dom.mem_actual > dom.last_target + XEN_FREE_MEM_LEFT / 4:
                 # VM didn't react to memory request at all, remove from donors.
                 if prev_mem_actual[domid] == dom.mem_actual:
@@ -371,7 +374,7 @@ def debug_stuck_balance(self, stuck_domid, memset_reqs, prev_mem_actual):
                     )
                     dom.slow_memset_react = True
 
-    def do_balance(self):
+    def do_balance(self) -> None:
         self.log.debug("do_balance()")
         if os.path.isfile("/var/run/qubes/do-not-membalance"):
             self.log.debug("do-not-membalance file present, returning")
@@ -388,7 +391,7 @@ def do_balance(self):
 
         self.print_stats(xenfree, memset_reqs)
 
-        prev_mem_actual = {}
+        prev_mem_actual: dict[str, Optional[int]] = {}
         for domid, dom in self.dom_dict.items():
             prev_mem_actual[domid] = dom.mem_actual
         for req in memset_reqs:
@@ -416,6 +419,7 @@ def do_balance(self):
                     self.debug_stuck_balance(
                         domid, memset_reqs, prev_mem_actual
                     )
+                    assert isinstance(dom.mem_actual, int)
                     self.mem_set(
                         domid,
                         self.get_free_xen_mem()