ominfra 0.0.0.dev101__py3-none-any.whl → 0.0.0.dev103__py3-none-any.whl

ominfra/scripts/journald2aws.py

@@ -4,38 +4,6 @@
 # @omlish-script
 # @omlish-amalg-output ../clouds/aws/journald2aws/main.py
 # ruff: noqa: N802 UP006 UP007 UP036
-"""
-TODO:
- - create log group
- - log stats - chunk sizes, byte count, num calls, etc
-
-==
-
-https://www.freedesktop.org/software/systemd/man/latest/journalctl.html
-
-journalctl:
-  -o json
-  --show-cursor
-
-  --since "2012-10-30 18:17:16"
-  --until "2012-10-30 18:17:16"
-
-  --after-cursor <cursor>
-
-==
-
-https://www.freedesktop.org/software/systemd/man/latest/systemd.journal-fields.html
-
-==
-
-@dc.dataclass(frozen=True)
-class Journald2AwsConfig:
-    log_group_name: str
-    log_stream_name: str
-
-    aws_batch_size: int = 1_000
-    aws_flush_interval_s: float = 1.
-"""
 import abc
 import argparse
 import base64
@@ -81,15 +49,21 @@ if sys.version_info < (3, 8):
 ########################################
 
 
-# ../../../../../omlish/lite/check.py
+# ../../../../../omlish/lite/cached.py
 T = ta.TypeVar('T')
 
+# ../../../../../omlish/lite/contextmanagers.py
+ExitStackedT = ta.TypeVar('ExitStackedT', bound='ExitStacked')
+
+# ../../../../threadworkers.py
+ThreadWorkerT = ta.TypeVar('ThreadWorkerT', bound='ThreadWorker')
+
 
 ########################################
 # ../../../../../omlish/lite/cached.py
 
 
-class cached_nullary:  # noqa
+class _cached_nullary:  # noqa
     def __init__(self, fn):
         super().__init__()
         self._fn = fn
@@ -106,6 +80,10 @@ class cached_nullary:  # noqa
         return bound
 
 
+def cached_nullary(fn: ta.Callable[..., T]) -> ta.Callable[..., T]:
+    return _cached_nullary(fn)
+
+
 ########################################
 # ../../../../../omlish/lite/check.py
 
@@ -222,7 +200,7 @@ class Pidfile:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        if self._f is not None:
+        if hasattr(self, '_f'):
             self._f.close()
             del self._f
 
@@ -736,6 +714,52 @@ class AwsDataclassMeta:
         return AwsDataclassMeta.Converters(d2a, a2d)
 
 
+########################################
+# ../../../../../omlish/lite/contextmanagers.py
+
+
+##
+
+
+class ExitStacked:
+    _exit_stack: ta.Optional[contextlib.ExitStack] = None
+
+    def __enter__(self: ExitStackedT) -> ExitStackedT:
+        check_state(self._exit_stack is None)
+        es = self._exit_stack = contextlib.ExitStack()
+        es.__enter__()
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        if (es := self._exit_stack) is None:
+            return None
+        return es.__exit__(exc_type, exc_val, exc_tb)
+
+    def _enter_context(self, cm: ta.ContextManager[T]) -> T:
+        es = check_not_none(self._exit_stack)
+        return es.enter_context(cm)
+
+
+##
+
+
+@contextlib.contextmanager
+def attr_setting(obj, attr, val, *, default=None):  # noqa
+    not_set = object()
+    orig = getattr(obj, attr, not_set)
+    try:
+        setattr(obj, attr, val)
+        if orig is not not_set:
+            yield orig
+        else:
+            yield default
+    finally:
+        if orig is not_set:
+            delattr(obj, attr)
+        else:
+            setattr(obj, attr, orig)
+
+
 ########################################
 # ../../../../../omlish/lite/io.py
 
@@ -1259,7 +1283,7 @@ class DataclassObjMarshaler(ObjMarshaler):
         return {k: m.marshal(getattr(o, k)) for k, m in self.fs.items()}
 
     def unmarshal(self, o: ta.Any) -> ta.Any:
-        return self.ty(**{k: self.fs[k].unmarshal(v) for k, v in o.items() if self.nonstrict or k in self.fs})
+        return self.ty(**{k: self.fs[k].unmarshal(v) for k, v in o.items() if not self.nonstrict or k in self.fs})
 
 
 @dc.dataclass(frozen=True)
@@ -1448,6 +1472,52 @@ def check_runtime_version() -> None:
             f'Requires python {REQUIRED_PYTHON_VERSION}, got {sys.version_info} from {sys.executable}')  # noqa
 
 
+########################################
+# ../cursor.py
+
+
+class JournalctlToAwsCursor:
+    def __init__(
+            self,
+            cursor_file: ta.Optional[str] = None,
+            *,
+            ensure_locked: ta.Optional[ta.Callable[[], None]] = None,
+    ) -> None:
+        super().__init__()
+        self._cursor_file = cursor_file
+        self._ensure_locked = ensure_locked
+
+    #
+
+    def get(self) -> ta.Optional[str]:
+        if self._ensure_locked is not None:
+            self._ensure_locked()
+
+        if not (cf := self._cursor_file):
+            return None
+        cf = os.path.expanduser(cf)
+
+        try:
+            with open(cf) as f:
+                return f.read().strip()
+        except FileNotFoundError:
+            return None
+
+    def set(self, cursor: str) -> None:
+        if self._ensure_locked is not None:
+            self._ensure_locked()
+
+        if not (cf := self._cursor_file):
+            return
+        cf = os.path.expanduser(cf)
+
+        log.info('Writing cursor file %s : %s', cf, cursor)
+        with open(ncf := cf + '.next', 'w') as f:
+            f.write(cursor)
+
+        os.rename(ncf, cf)
+
+
 ########################################
 # ../../logs.py
 """
@@ -1502,7 +1572,7 @@ class AwsPutLogEventsResponse(AwsDataclass):
 ##
 
 
-class AwsLogMessagePoster:
+class AwsLogMessageBuilder:
     """
     TODO:
      - max_items
@@ -1528,7 +1598,7 @@ class AwsLogMessagePoster:
             log_group_name: str,
             log_stream_name: str,
             region_name: str,
-            credentials: AwsSigner.Credentials,
+            credentials: ta.Optional[AwsSigner.Credentials],
 
             url: ta.Optional[str] = None,
             service_name: str = DEFAULT_SERVICE_NAME,
@@ -1550,11 +1620,16 @@ class AwsLogMessagePoster:
             headers = {**headers, **extra_headers}
         self._headers = {k: [v] for k, v in headers.items()}
 
-        self._signer = V4AwsSigner(
-            credentials,
-            region_name,
-            service_name,
-        )
+        signer: ta.Optional[V4AwsSigner]
+        if credentials is not None:
+            signer = V4AwsSigner(
+                credentials,
+                region_name,
+                service_name,
+            )
+        else:
+            signer = None
+        self._signer = signer
 
     #
 
@@ -1598,13 +1673,14 @@ class AwsLogMessagePoster:
             payload=body,
         )
 
-        sig_headers = self._signer.sign(
-            sig_req,
-            sign_payload=False,
-        )
-        sig_req = dc.replace(sig_req, headers={**sig_req.headers, **sig_headers})
+        if (signer := self._signer) is not None:
+            sig_headers = signer.sign(
+                sig_req,
+                sign_payload=False,
+            )
+            sig_req = dc.replace(sig_req, headers={**sig_req.headers, **sig_headers})
 
-        post = AwsLogMessagePoster.Post(
+        post = AwsLogMessageBuilder.Post(
             url=self._url,
             headers={k: check_single(v) for k, v in sig_req.headers.items()},
             data=sig_req.payload,
@@ -1687,15 +1763,20 @@ class JournalctlMessageBuilder:
 
 
 ########################################
-# ../../../../threadworker.py
+# ../../../../threadworkers.py
 """
 TODO:
  - implement stop lol
  - collective heartbeat monitoring - ThreadWorkerGroups
+ - group -> 'context'? :|
+  - shared stop_event?
 """
 
 
-class ThreadWorker(abc.ABC):
+##
+
+
+class ThreadWorker(ExitStacked, abc.ABC):
     def __init__(
             self,
             *,
@@ -1707,46 +1788,107 @@ class ThreadWorker(abc.ABC):
             stop_event = threading.Event()
         self._stop_event = stop_event
 
+        self._lock = threading.RLock()
         self._thread: ta.Optional[threading.Thread] = None
-
         self._last_heartbeat: ta.Optional[float] = None
 
     #
 
+    def __enter__(self: ThreadWorkerT) -> ThreadWorkerT:
+        with self._lock:
+            return super().__enter__()  # noqa
+
+    #
+
     def should_stop(self) -> bool:
         return self._stop_event.is_set()
 
+    class Stopping(Exception):  # noqa
+        pass
+
     #
 
     @property
     def last_heartbeat(self) -> ta.Optional[float]:
         return self._last_heartbeat
 
-    def _heartbeat(self) -> bool:
+    def _heartbeat(
+            self,
+            *,
+            no_stop_check: bool = False,
+    ) -> None:
         self._last_heartbeat = time.time()
 
-        if self.should_stop():
+        if not no_stop_check and self.should_stop():
             log.info('Stopping: %s', self)
-            return False
-
-        return True
+            raise ThreadWorker.Stopping
 
     #
 
+    def has_started(self) -> bool:
+        return self._thread is not None
+
     def is_alive(self) -> bool:
         return (thr := self._thread) is not None and thr.is_alive()
 
     def start(self) -> None:
-        thr = threading.Thread(target=self._run)
-        self._thread = thr
-        thr.start()
+        with self._lock:
+            if self._thread is not None:
+                raise RuntimeError('Thread already started: %r', self)
+
+            thr = threading.Thread(target=self.__run)
+            self._thread = thr
+            thr.start()
+
+    #
+
+    def __run(self) -> None:
+        try:
+            self._run()
+        except ThreadWorker.Stopping:
+            log.exception('Thread worker stopped: %r', self)
+        except Exception:  # noqa
+            log.exception('Error in worker thread: %r', self)
+            raise
 
     @abc.abstractmethod
     def _run(self) -> None:
         raise NotImplementedError
 
+    #
+
     def stop(self) -> None:
-        raise NotImplementedError
+        self._stop_event.set()
+
+    def join(self, timeout: ta.Optional[float] = None) -> None:
+        with self._lock:
+            if self._thread is None:
+                raise RuntimeError('Thread not started: %r', self)
+            self._thread.join(timeout)
+
+
+##
+
+
+class ThreadWorkerGroup:
+    @dc.dataclass()
+    class State:
+        worker: ThreadWorker
+
+    def __init__(self) -> None:
+        super().__init__()
+
+        self._lock = threading.RLock()
+        self._states: ta.Dict[ThreadWorker, ThreadWorkerGroup.State] = {}
+
+    def add(self, *workers: ThreadWorker) -> 'ThreadWorkerGroup':
+        with self._lock:
+            for w in workers:
+                if w in self._states:
+                    raise KeyError(w)
+                self._states[w] = ThreadWorkerGroup.State(w)
+
+        return self
 
 
 ########################################
@@ -1855,6 +1997,103 @@ def subprocess_try_output_str(*args: str, **kwargs: ta.Any) -> ta.Optional[str]:
     return out.decode().strip() if out is not None else None
 
 
+##
+
+
+def subprocess_close(
+        proc: subprocess.Popen,
+        timeout: ta.Optional[float] = None,
+) -> None:
+    # TODO: terminate, sleep, kill
+    if proc.stdout:
+        proc.stdout.close()
+    if proc.stderr:
+        proc.stderr.close()
+    if proc.stdin:
+        proc.stdin.close()
+
+    proc.wait(timeout)
+
+
+########################################
+# ../poster.py
+"""
+TODO:
+ - retries
+"""
+
+
+class JournalctlToAwsPosterWorker(ThreadWorker):
+    def __init__(
+            self,
+            queue,  # type: queue.Queue[ta.Sequence[JournalctlMessage]]  # noqa
+            builder: AwsLogMessageBuilder,
+            cursor: JournalctlToAwsCursor,
+            *,
+            ensure_locked: ta.Optional[ta.Callable[[], None]] = None,
+            dry_run: bool = False,
+            queue_timeout_s: float = 1.,
+            **kwargs: ta.Any,
+    ) -> None:
+        super().__init__(**kwargs)
+        self._queue = queue
+        self._builder = builder
+        self._cursor = cursor
+        self._ensure_locked = ensure_locked
+        self._dry_run = dry_run
+        self._queue_timeout_s = queue_timeout_s
+    #
+
+    def _run(self) -> None:
+        if self._ensure_locked is not None:
+            self._ensure_locked()
+
+        last_cursor: ta.Optional[str] = None  # noqa
+        while True:
+            self._heartbeat()
+
+            try:
+                msgs: ta.Sequence[JournalctlMessage] = self._queue.get(timeout=self._queue_timeout_s)
+            except queue.Empty:
+                msgs = []
+
+            if not msgs:
+                log.debug('Empty queue chunk')
+                continue
+
+            log.debug('%r', msgs)
+
+            cur_cursor: ta.Optional[str] = None
+            for m in reversed(msgs):
+                if m.cursor is not None:
+                    cur_cursor = m.cursor
+                    break
+
+            feed_msgs = []
+            for m in msgs:
+                feed_msgs.append(AwsLogMessageBuilder.Message(
+                    message=json.dumps(m.dct, sort_keys=True),
+                    ts_ms=int((m.ts_us / 1000.) if m.ts_us is not None else (time.time() * 1000.)),
+                ))
+
+            for post in self._builder.feed(feed_msgs):
+                log.debug('%r', post)
+
+                if not self._dry_run:
+                    with urllib.request.urlopen(urllib.request.Request(  # noqa
+                            post.url,
+                            method='POST',
+                            headers=dict(post.headers),
+                            data=post.data,
+                    )) as resp:
+                        response = AwsPutLogEventsResponse.from_aws(json.loads(resp.read().decode('utf-8')))
+                    log.debug('%r', response)
+
+            if cur_cursor is not None:
+                self._cursor.set(cur_cursor)
+                last_cursor = cur_cursor  # noqa
+
+
 ########################################
 # ../../../../journald/tailer.py
 """
@@ -2227,7 +2466,7 @@ class JournalctlTailerWorker(ThreadWorker):
         self._read_size = read_size
         self._sleep_s = sleep_s
 
-        self._mb = JournalctlMessageBuilder()
+        self._builder = JournalctlMessageBuilder()
 
         self._proc: ta.Optional[subprocess.Popen] = None
 
@@ -2251,69 +2490,103 @@ class JournalctlTailerWorker(ThreadWorker):
 
         return cmd
 
+    def _read_loop(self, stdout: ta.IO) -> None:
+        while stdout.readable():
+            self._heartbeat()
+
+            buf = stdout.read(self._read_size)
+            if not buf:
+                log.debug('Journalctl empty read')
+                break
+
+            log.debug('Journalctl read buffer: %r', buf)
+            msgs = self._builder.feed(buf)
+            if msgs:
+                while True:
+                    try:
+                        self._output.put(msgs, timeout=1.)
+                    except queue.Full:
+                        self._heartbeat()
+                    else:
+                        break
+
     def _run(self) -> None:
         with subprocess.Popen(
             self._full_cmd(),
             stdout=subprocess.PIPE,
         ) as self._proc:
-            stdout = check_not_none(self._proc.stdout)
+            try:
+                stdout = check_not_none(self._proc.stdout)
 
-            fd = stdout.fileno()
-            fl = fcntl.fcntl(fd, fcntl.F_GETFL)
-            fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
+                fd = stdout.fileno()
+                fl = fcntl.fcntl(fd, fcntl.F_GETFL)
+                fcntl.fcntl(fd, fcntl.F_SETFL, fl | os.O_NONBLOCK)
 
-            while True:
-                if not self._heartbeat():
-                    return
+                while True:
+                    self._heartbeat()
 
-                while stdout.readable():
-                    if not self._heartbeat():
-                        return
+                    self._read_loop(stdout)
 
-                    buf = stdout.read(self._read_size)
-                    if not buf:
-                        log.debug('Journalctl empty read')
-                        break
+                    log.debug('Journalctl not readable')
 
-                    log.debug('Journalctl read buffer: %r', buf)
-                    msgs = self._mb.feed(buf)
-                    if msgs:
-                        while True:
-                            try:
-                                self._output.put(msgs, timeout=1.)
-                            except queue.Full:
-                                if not self._heartbeat():
-                                    return
-                            else:
-                                break
+                    if self._proc.poll() is not None:
+                        log.critical('Journalctl process terminated')
+                        return
 
-                if self._proc.poll() is not None:
-                    log.critical('Journalctl process terminated')
-                    return
+                    time.sleep(self._sleep_s)
 
-                log.debug('Journalctl readable')
-                time.sleep(self._sleep_s)
+            finally:
+                subprocess_close(self._proc)
 
 
 ########################################
-# main.py
+# ../driver.py
+"""
+TODO:
+ - create log group
+ - log stats - chunk sizes, byte count, num calls, etc
+
+==
 
+https://www.freedesktop.org/software/systemd/man/latest/journalctl.html
+
+journalctl:
+  -o json
+  --show-cursor
+
+  --since "2012-10-30 18:17:16"
+  --until "2012-10-30 18:17:16"
+
+  --after-cursor <cursor>
+
+==
+
+https://www.freedesktop.org/software/systemd/man/latest/systemd.journal-fields.html
+
+==
 
 @dc.dataclass(frozen=True)
-class JournalctlOpts:
-    after_cursor: ta.Optional[str] = None
+class Journald2AwsConfig:
+    log_group_name: str
+    log_stream_name: str
 
-    since: ta.Optional[str] = None
-    until: ta.Optional[str] = None
+    aws_batch_size: int = 1_000
+    aws_flush_interval_s: float = 1.
+"""
+
+
+##
 
 
-class JournalctlToAws:
+class JournalctlToAwsDriver(ExitStacked):
     @dc.dataclass(frozen=True)
     class Config:
         pid_file: ta.Optional[str] = None
 
         cursor_file: ta.Optional[str] = None
 
+        runtime_limit: ta.Optional[float] = None
+
         #
 
         aws_log_group_name: str = 'omlish'
@@ -2324,6 +2597,8 @@ class JournalctlToAws:
 
         aws_region_name: str = 'us-west-1'
 
+        aws_dry_run: bool = False
+
         #
 
         journalctl_cmd: ta.Optional[ta.Sequence[str]] = None
@@ -2331,24 +2606,10 @@ class JournalctlToAws:
         journalctl_after_cursor: ta.Optional[str] = None
         journalctl_since: ta.Optional[str] = None
 
-        #
-
-        dry_run: bool = False
-
     def __init__(self, config: Config) -> None:
         super().__init__()
-        self._config = config
-
-    #
-
-    _es: contextlib.ExitStack
 
-    def __enter__(self) -> 'JournalctlToAws':
-        self._es = contextlib.ExitStack().__enter__()
-        return self
-
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        return self._es.__exit__(exc_type, exc_val, exc_tb)
+        self._config = config
 
     #
 
@@ -2361,7 +2622,7 @@ class JournalctlToAws:
 
         log.info('Opening pidfile %s', pfp)
 
-        pf = self._es.enter_context(Pidfile(pfp))
+        pf = self._enter_context(Pidfile(pfp))
         pf.write()
         return pf
 
@@ -2371,48 +2632,32 @@ class JournalctlToAws:
 
     #
 
-    def _read_cursor_file(self) -> ta.Optional[str]:
-        self._ensure_locked()
-
-        if not (cf := self._config.cursor_file):
-            return None
-        cf = os.path.expanduser(cf)
-
-        try:
-            with open(cf) as f:
-                return f.read().strip()
-        except FileNotFoundError:
-            return None
-
-    def _write_cursor_file(self, cursor: str) -> None:
-        self._ensure_locked()
-
-        if not (cf := self._config.cursor_file):
-            return
-        cf = os.path.expanduser(cf)
-
-        log.info('Writing cursor file %s : %s', cf, cursor)
-        with open(ncf := cf + '.next', 'w') as f:
-            f.write(cursor)
-
-        os.rename(ncf, cf)
+    @cached_nullary
+    def _cursor(self) -> JournalctlToAwsCursor:
+        return JournalctlToAwsCursor(
+            self._config.cursor_file,
+            ensure_locked=self._ensure_locked,
+        )
 
     #
 
     @cached_nullary
-    def _aws_credentials(self) -> AwsSigner.Credentials:
+    def _aws_credentials(self) -> ta.Optional[AwsSigner.Credentials]:
+        if self._config.aws_access_key_id is None and self._config.aws_secret_access_key is None:
+            return None
+
         return AwsSigner.Credentials(
             access_key_id=check_non_empty_str(self._config.aws_access_key_id),
             secret_access_key=check_non_empty_str(self._config.aws_secret_access_key),
         )
 
     @cached_nullary
-    def _aws_log_message_poster(self) -> AwsLogMessagePoster:
-        return AwsLogMessagePoster(
+    def _aws_log_message_builder(self) -> AwsLogMessageBuilder:
+        return AwsLogMessageBuilder(
             log_group_name=self._config.aws_log_group_name,
             log_stream_name=check_non_empty_str(self._config.aws_log_stream_name),
             region_name=self._config.aws_region_name,
-            credentials=check_not_none(self._aws_credentials()),
+            credentials=self._aws_credentials(),
         )
 
     #
@@ -2431,7 +2676,7 @@ class JournalctlToAws:
         else:
             ac = self._config.journalctl_after_cursor
             if ac is None:
-                ac = self._read_cursor_file()
+                ac = self._cursor().get()
             if ac is not None:
                 log.info('Starting from cursor %s', ac)
 
@@ -2447,63 +2692,49 @@ class JournalctlToAws:
 
     #
 
-    def run(self) -> None:
-        self._ensure_locked()
+    @cached_nullary
+    def _aws_poster_worker(self) -> JournalctlToAwsPosterWorker:
+        return JournalctlToAwsPosterWorker(
+            self._journalctl_message_queue(),
+            self._aws_log_message_builder(),
+            self._cursor(),
 
-        q = self._journalctl_message_queue()  # type: queue.Queue[ta.Sequence[JournalctlMessage]]
-        jtw = self._journalctl_tailer_worker()  # type: JournalctlTailerWorker
-        mp = self._aws_log_message_poster()  # type: AwsLogMessagePoster
+            ensure_locked=self._ensure_locked,
+            dry_run=self._config.aws_dry_run,
+        )
 
-        jtw.start()
+    #
 
-        last_cursor: ta.Optional[str] = None  # noqa
-        while True:
-            if not jtw.is_alive():
-                log.critical('Journalctl tailer worker died')
-                break
+    def run(self) -> None:
+        pw: JournalctlToAwsPosterWorker = self._aws_poster_worker()
+        tw: JournalctlTailerWorker = self._journalctl_tailer_worker()
 
-            try:
-                msgs: ta.Sequence[JournalctlMessage] = q.get(timeout=1.)
-            except queue.Empty:
-                msgs = []
-            if not msgs:
-                continue
+        ws = [pw, tw]
 
-            log.debug('%r', msgs)
+        for w in ws:
+            w.start()
 
-            cur_cursor: ta.Optional[str] = None
-            for m in reversed(msgs):
-                if m.cursor is not None:
-                    cur_cursor = m.cursor
+        start = time.time()
+
+        while True:
+            for w in ws:
+                if not w.is_alive():
+                    log.critical('Worker died: %r', w)
                     break
 
-            if not msgs:
-                log.warning('Empty queue chunk')
-                continue
+            if (rl := self._config.runtime_limit) is not None and time.time() - start >= rl:
+                log.warning('Runtime limit reached')
+                break
 
-            feed_msgs = []
-            for m in msgs:
-                feed_msgs.append(mp.Message(
-                    message=json.dumps(m.dct, sort_keys=True),
-                    ts_ms=int((m.ts_us / 1000.) if m.ts_us is not None else (time.time() * 1000.)),
-                ))
+            time.sleep(1.)
 
-            [post] = mp.feed(feed_msgs)
-            log.debug('%r', post)
+        for w in reversed(ws):
+            w.stop()
+            w.join()
 
-            if not self._config.dry_run:
-                with urllib.request.urlopen(urllib.request.Request(  # noqa
-                        post.url,
-                        method='POST',
-                        headers=dict(post.headers),
-                        data=post.data,
-                )) as resp:
-                    response = AwsPutLogEventsResponse.from_aws(json.loads(resp.read().decode('utf-8')))
-                log.debug('%r', response)
 
-            if cur_cursor is not None:
-                self._write_cursor_file(cur_cursor)
-                last_cursor = cur_cursor  # noqa
+########################################
+# main.py
 
 
 def _main() -> None:
@@ -2518,6 +2749,8 @@ def _main() -> None:
 
     parser.add_argument('--message', nargs='?')
     parser.add_argument('--real', action='store_true')
+    parser.add_argument('--num-messages', type=int)
+    parser.add_argument('--runtime-limit', type=float)
 
     args = parser.parse_args()
 
@@ -2527,13 +2760,13 @@ def _main() -> None:
 
     #
 
-    config: JournalctlToAws.Config
+    config: JournalctlToAwsDriver.Config
     if args.config_file:
         with open(os.path.expanduser(args.config_file)) as cf:
             config_dct = json.load(cf)
-        config = unmarshal_obj(config_dct, JournalctlToAws.Config)
+        config = unmarshal_obj(config_dct, JournalctlToAwsDriver.Config)
     else:
-        config = JournalctlToAws.Config()
+        config = JournalctlToAwsDriver.Config()
 
     #
 
@@ -2550,7 +2783,7 @@ def _main() -> None:
             '--sleep-n', '2',
             '--sleep-s', '.5',
             *(['--message', args.message] if args.message else []),
-            '100000',
+            str(args.num_messages or 100_000),
         ])
 
     #
@@ -2558,14 +2791,14 @@ def _main() -> None:
     for ca, pa in [
         ('journalctl_after_cursor', 'after_cursor'),
         ('journalctl_since', 'since'),
-        ('dry_run', 'dry_run'),
+        ('aws_dry_run', 'dry_run'),
     ]:
         if (av := getattr(args, pa)):
             config = dc.replace(config, **{ca: av})
 
     #
 
-    with JournalctlToAws(config) as jta:
+    with JournalctlToAwsDriver(config) as jta:
         jta.run()