mosaic.py

#!/usr/bin/env python3 # Mosaic Emulator and Checker import argparse import ast import copy import inspect import json import random from dataclasses import dataclass from itertools import compress, product from pathlib import Path from typing import Callable, Generator ## 1. Mosaic system calls ### 1.1 Process, thread, and context switching sys_fork = lambda: os.sys_fork() sys_spawn = lambda fn, *args: os.sys_spawn(fn, *args) sys_sched = lambda: os.sys_sched() ### 1.2 Virtual character device sys_choose = lambda choices: os.sys_choose(choices) sys_write = lambda *args: os.sys_write(*args) ### 1.3 Virtual block storage device sys_bread = lambda k: os.sys_bread(k) sys_bwrite = lambda k, v: os.sys_bwrite(k, v) sys_sync = lambda: os.sys_sync() sys_crash = lambda: os.sys_crash() ### 1.4 System call helpers SYSCALLS = [] def syscall(func): # @syscall decorator SYSCALLS.append(func.__name__) return func ## 2. Mosaic operating system emulator ### 2.1 Data structures class Heap: pass # no member: self.__dict__ is the heap @dataclass class Thread: context: Generator # program counter, local variables, etc. heap: Heap # a pointer to thread's "memory" @dataclass class Storage: persist: dict # persisted storage state buf: dict # outstanding operations yet to be persisted ### 2.2 The OperatingSystem class class OperatingSystem: """An executable operating system model. The operating system model hosts a single Python application with a main() function accessible to a shared heap and 9 system calls (marked by the @syscall decorator). An example: def main(): pid = sys_fork() sys_sched() # non-deterministic context switch if pid == 0: sys_write('World') else: sys_write('Hello') At any moment, this model keeps tracking a set of threads and a "currently running" one. Each thread consists of a reference to a heap object (may be shared with other threads) and a private context (program counters, local variables, etc.). A thread context is a Python generator object, i.e., a stack-less coroutine [1] that can save the running context and yield itself. For applications, the keyword "yield" is reserved for system calls. For example, a "choose" system call [2]: sys_choose(['A', 'B']) is transpiled as yielding the string "sys_choose" followed by its parameters (choices): res = yield 'sys_choose', ['A', 'B']. Yield will transfer the control to the OS for system call handling and eventually returning a value ('A' or 'B') to the application. Right after transferring control to the OS by "yield", the function state is "frozen", where program counters and local variables are accessible via the generator object. Therefore, OS can serialize its internal state--all thread's contexts, heaps, and virtual device states at this moment. In this sense, operating system is a system-call driven state transition system: (s0) --run first thread (main)-> (s1) --sys_choose and application execution-> (s2) --sys_sched and application execution-> (s3) ... Real operating systems can be preemptive--context switching can happen non-deterministically at any program point, simply because processor can non-deterministically interrupt its currently running code and transfer the control to the operating system. The OS internal implementation does NOT immediately process the system call: it returns all possible choices available at the moment and their corresponding processing logic as callbacks. For the example above, the "choose" system call returns a non-deterministic choice among given choices. The internal implementation thus returns choices = { 'A': (lambda: 'A'), 'B': (lambda: 'B'), } for later processing. Another example is non-deterministic context switching by yielding 'sys_sched'. Suppose there are threads t1 and t2 at the moment. The system call handler will return choices = { 't1': (lambda: switch_to(t1)), 't2': (lambda: switch_to(t2)), } in which switch_to(th) replaces the OS's current running thread with th (changes the global "heap" variable). Such deferred execution of system calls separates the mechanism of non-deterministic choices from the actual decision makers (e.g., an interpreter or a model checker). Once the decision is made, the simply call step(choice) and the OS will execute this choice by choices[choice]() with the application code (generator) being resumed. This model provides "write" system call to immediately push data to a hypothetical character device like a tty associated with stdout. We model a block device (key-value store) that may lose data upon crash. The model assumes atomicity of each single block write (a key-value update). However, writes are merely to a volatile buffer which may non-deterministically lose data upon crash 3]. The "sync" system call persists buffered writes. References: [1] N. Schemenauer, T. Peters, and M. L. Hetland. PEP 255 - Simple generators. https://peps.python.org/pep-0255/ [2] J. Yang, C. Sar, and D. Engler. eXplode: a lightweight, general system for finding serious storage system errors. OSDI'06. [3] T. S. Pillai, V. Chidambaram, R. Alagappan, A. Al-Kiswany, A. C. Arpaci-Dusseau, and R. H. Arpaci-Dusseau. All file systems are not created equal: On the complexity of crafting crash consistent applications. OSDI'14. """ def __init__(self, init: Callable): """Create a new OS instance with pending-to-execute init thread.""" # Operating system states self._threads = [Thread(context=init(), heap=Heap())] self._current = 0 self._choices = {init.__name__: lambda: None} self._stdout = '' self._storage = Storage(persist={}, buf={}) # Internal states self._init = init self._trace = [] self._newfork = set() ### 2.3 System call implementation #### 2.3.1 Process, thread, and context switching @syscall def sys_spawn(self, func: Callable, *args): """Spawn a heap-sharing new thread executing func(args).""" def do_spawn(): self._threads.append( Thread( context=func(*args), # func() returns a new generator heap=self.current().heap, # shared heap ) ) return {'spawn': (lambda: do_spawn())} @syscall def sys_fork(self): """Create a clone of the current thread with a copied heap.""" if all(not f.frame.f_locals['fork_child'] for f in inspect.stack() if f.function == '_step'): # this is parent; do fork # Deep-copying generators causes troubles--they are twined with # Python's runtime state. We use an (inefficient) hack here: replay # the entire trace and override the last fork() to avoid infinite # recursion. os_clone = OperatingSystem(self._init) os_clone.replay(self._trace[:-1]) os_clone._step(self._trace[-1], fork_child=True) # Now os_clone._current is the forked process. Cloned thread just # yields a sys_fork and is pending for fork()'s return value. It # is necessary to mark cloned threads (in self._newfork) and send # child's fork() return value when they are scheduled for the # first time. def do_fork(): self._threads.append(os_clone.current()) self._newfork.add((pid := len(self._threads)) - 1) return 1000 + pid # returned pid starts from 1000 return {'fork': (lambda: do_fork())} else: return None # overridden fork; this value is never used because # os_clone is destroyed immediately after fork() @syscall def sys_sched(self): """Return a non-deterministic context switch to a runnable thread.""" return { f't{i+1}': (lambda i=i: self._switch_to(i)) for i, th in enumerate(self._threads) if th.context.gi_frame is not None # thread still alive? } ### 2.3.2 Virtual character device (byte stream) @syscall def sys_choose(self, choices): """Return a non-deterministic value from choices.""" return {f'choose {c}': (lambda c=c: c) for c in choices} @syscall def sys_write(self, *args): """Write strings (space separated) to stdout.""" def do_write(): self._stdout += ' '.join(str(arg) for arg in args) return {'write': (lambda: do_write())} ### 2.3.3 Virtual block storage device @syscall def sys_bread(self, key): """Return the specific key's associated value in block device.""" storage = self._storage return {'bread': (lambda: storage.buf.get(key, # always try to read from buffer first storage.persist.get(key, None) # and then persistent storage ) )} @syscall def sys_bwrite(self, key, value): """Write (key, value) pair to block device's buffer.""" def do_bwrite(): self._storage.buf[key] = value return {'bwrite': (lambda: do_bwrite())} @syscall def sys_sync(self): """Persist all buffered writes.""" def do_sync(): store = self._storage self._storage = Storage( persist=store.persist | store.buf, # write back buf={} ) return {'sync': (lambda: do_sync())} @syscall def sys_crash(self): """Simulate a system crash that non-deterministically persists outstanding writes in the buffer. """ persist = self._storage.persist btrace = self._storage.buf.items() # block trace crash_sites = ( lambda subset=subset: setattr(self, '_storage', Storage( # persist only writes in the subset persist=persist | dict(compress(btrace, subset)), buf={} ) ) for subset in # Mosaic allows persisting any subset of product( # pending blocks in the buffer *([(0, 1)] * len(btrace)) ) ) return dict(enumerate(crash_sites)) ### 2.4 Operating system as a state machine def replay(self, trace: list) -> dict: """Replay an execution trace and return the resulting state.""" for choice in trace: self._step(choice) return self.state_dump() def _step(self, choice, fork_child=False): self._switch_to(self._current) self._trace.append(choice) # keep all choices for replay-based fork() action = self._choices[choice] # return value of sys_xxx: a lambda res = action() try: # Execute current thread for one step func, args = self.current().context.send(res) assert func in SYSCALLS self._choices = getattr(self, func)(*args) except StopIteration: # ... and thread terminates self._choices = self.sys_sched() # At this point, the operating system's state is # (self._threads, self._current, self._stdout, self._storage) # and outgoing transitions are saved in self._choices. ### 2.5 Misc and helper functions def state_dump(self) -> dict: """Create a serializable Mosaic state dump with hash code.""" heaps = {} for th in self._threads: if (i := id(th.heap)) not in heaps: # unique heaps heaps[i] = len(heaps) + 1 os_state = { 'current': self._current, 'choices': sorted(list(self._choices.keys())), 'contexts': [ { 'name': th.context.gi_frame.f_code.co_name, 'heap': heaps[id(th.heap)], # the unique heap id 'pc': th.context.gi_frame.f_lineno, 'locals': th.context.gi_frame.f_locals, } if th.context.gi_frame is not None else None for th in self._threads ], 'heaps': { heaps[id(th.heap)]: th.heap.__dict__ for th in self._threads }, 'stdout': self._stdout, 'store_persist': self._storage.persist, 'store_buffer': self._storage.buf, } h = hash(json.dumps(os_state, sort_keys=True)) + 2**63 return (copy.deepcopy(os_state) # freeze the runtime state | dict(hashcode=f'{h:016x}')) def current(self) -> Thread: """Return the current running thread object.""" return self._threads[self._current] def _switch_to(self, tid: int): self._current = tid globals()['os'] = self globals()['heap'] = self.current().heap if tid in self._newfork: self._newfork.remove(tid) # tricky: forked process must receive 0 return 0 # to indicate a child ## 3. The Mosaic runtime class Mosaic: """The operating system interpreter and model checker. The operating system model is a state transition system: os.replay() maps any trace to a state (with its outgoing transitions). Based on this model, two state space explorers are implemented: - run: Choose outgoing transitions uniformly at random, yielding a single execution trace. - check: Exhaustively explore all reachable states by a breadth- first search. Duplicated states are not visited twice. Both explorers produce the visited portion of the state space as a serializable object containing: - source: The application source code - vertices: A list of operating system state dumps. The first vertex in the list is the initial state. Each vertex has a unique "hashcode" id. - edges: A list of 3-tuples: (source, target, label) denoting an explored source --[label]-> target edge. Both source and target are state hashcode ids. """ ### 3.1 Model interpreter and checker def run(self) -> dict: """Interpret the model with non-deterministic choices.""" os = OperatingSystem(self.entry) V, E = [os.state_dump() | dict(depth=0)], [] while (choices := V[-1]['choices']): choice = random.choice(choices) # uniformly at random V.append(os.replay([choice]) | dict(depth=len(V))) E.append((V[-2]['hashcode'], V[-1]['hashcode'], choice)) return dict(source=self.src, vertices=V, edges=E) def check(self) -> dict: """Exhaustively explore the state space.""" class State: entry = self.entry def __init__(self, trace): self.trace = trace self.state = OperatingSystem(State.entry).replay(trace) self.state |= dict(depth=0) self.hashcode = self.state['hashcode'] def extend(self, c): st = State(self.trace + (c,)) st.state = st.state | dict(depth=self.state['depth'] + 1) return st st0 = State(tuple()) # initial state of empty trace queued, V, E = [st0], {st0.hashcode: st0.state}, [] while queued: st = queued.pop(0) for choice in st.state['choices']: st1 = st.extend(choice) if st1.hashcode not in V: # found an unexplored state V[st1.hashcode] = st1.state queued.append(st1) E.append((st.hashcode, st1.hashcode, choice)) return dict( source=self.src, vertices=sorted(V.values(), key=lambda st: st['depth']), edges=E ) ### 3.1 Source code parsing and rewriting class Transformer(ast.NodeTransformer): def visit_Call(self, node): # Rewrite system calls as yields if (isinstance(node.func, ast.Name) and node.func.id in SYSCALLS): # rewrite system calls return ast.Yield(ast.Tuple( # -> yield ('sys_xxx', args) elts=[ ast.Constant(value=node.func.id), ast.Tuple(elts=node.args), ] )) else: return node def __init__(self, src: str): tree = ast.parse(src) hacked_ast = self.Transformer().visit(tree) hacked_src = ast.unparse(hacked_ast) context = {} exec(hacked_src, globals(), context) globals().update(context) self.src = src self.entry = context['main'] # must have a main() ## 4. Utilities if __name__ == '__main__': parser = argparse.ArgumentParser( description='The modeled operating system and state explorer.' ) parser.add_argument( 'source', help='application code (.py) to be checked; must have a main()' ) parser.add_argument('-r', '--run', action='store_true') parser.add_argument('-c', '--check', action='store_true') args = parser.parse_args() src = Path(args.source).read_text() mosaic = Mosaic(src) if args.check: explored = mosaic.check() else: explored = mosaic.run() # run is the default option # Serialize the explored states and write to stdout. This encourages piping # the results to another tool following the UNIX philosophy. Examples: # # mosaic --run foo.py | grep stdout | tail -n 1 # quick and dirty check # mosaic --check bar.py | fx # or any other interactive visualizer # print(json.dumps(explored, ensure_ascii=False, indent=2))