Skip to content
rachel/workerd
File

Blob: src/workerd/util/state-machine-test.c++

28.2 KB
1// Copyright (c) 2017-2025 Cloudflare, Inc.
2// Licensed under the Apache 2.0 license found in the LICENSE file or at:
3//     https://opensource.org/licenses/Apache-2.0
4 
5#include "state-machine.h"
6 
7#include <kj/test.h>
8 
9// Entire test file was Claude-generated initially.
10 
11namespace workerd {
12namespace {
13 
14// =============================================================================
15// Test State Types
16// =============================================================================
17 
18struct Idle {
19  static constexpr kj::StringPtr NAME KJ_UNUSED = "idle"_kj;
20  bool initialized = false;
21};
22 
23struct Running {
24  static constexpr kj::StringPtr NAME KJ_UNUSED = "running"_kj;
25  kj::String taskName;
26  int progress = 0;
27 
28  Running() = default;
29  explicit Running(kj::String name): taskName(kj::mv(name)) {}
30};
31 
32struct Completed {
33  static constexpr kj::StringPtr NAME KJ_UNUSED = "completed"_kj;
34  int result;
35 
36  explicit Completed(int r): result(r) {}
37};
38 
39struct Failed {
40  static constexpr kj::StringPtr NAME KJ_UNUSED = "failed"_kj;
41  kj::String error;
42 
43  explicit Failed(kj::String err): error(kj::mv(err)) {}
44};
45 
46// =============================================================================
47// Basic StateMachine Tests
48// =============================================================================
49 
50KJ_TEST("StateMachine: basic state checks") {
51  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Idle>();
52 
53  // Initialized to Idle via create()
54  KJ_EXPECT(machine.isInitialized());
55  KJ_EXPECT(machine.is<Idle>());
56  KJ_EXPECT(!machine.is<Running>());
57}
58 
59KJ_TEST("StateMachine: state data access") {
60  auto machine =
61      StateMachine<Idle, Running, Completed, Failed>::create<Running>(kj::str("my-task"));
62 
63  KJ_EXPECT(machine.is<Running>());
64  auto& running = machine.getUnsafe<Running>();
65  KJ_EXPECT(running.taskName == "my-task");
66  KJ_EXPECT(running.progress == 0);
67 
68  // Modify state data
69  running.progress = 50;
70  KJ_EXPECT(machine.getUnsafe<Running>().progress == 50);
71}
72 
73KJ_TEST("StateMachine: tryGet returns none for wrong state") {
74  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Idle>();
75 
76  // tryGet for correct state
77  KJ_IF_SOME(idle, machine.tryGetUnsafe<Idle>()) {
78    KJ_EXPECT(!idle.initialized);
79  } else {
80    KJ_FAIL_EXPECT("Should have gotten Idle state");
81  }
82 
83  // tryGet for wrong state
84  KJ_EXPECT(machine.tryGetUnsafe<Running>() == kj::none);
85  KJ_EXPECT(machine.tryGetUnsafe<Completed>() == kj::none);
86}
87 
88KJ_TEST("StateMachine: isAnyOf checks multiple states") {
89  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Completed>(42);
90 
91  // Use local variables to avoid KJ_EXPECT macro parsing issues with template brackets
92  bool isCompletedOrFailed = machine.isAnyOf<Completed, Failed>();
93  bool isIdleOrRunning = machine.isAnyOf<Idle, Running>();
94  KJ_EXPECT(isCompletedOrFailed);
95  KJ_EXPECT(!isIdleOrRunning);
96 
97  machine.transitionTo<Failed>(kj::str("error"));
98  isCompletedOrFailed = machine.isAnyOf<Completed, Failed>();
99  isIdleOrRunning = machine.isAnyOf<Idle, Running>();
100  KJ_EXPECT(isCompletedOrFailed);
101  KJ_EXPECT(!isIdleOrRunning);
102}
103 
104KJ_TEST("StateMachine: transitionFromTo with precondition") {
105  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Idle>();
106 
107  // Transition from wrong state fails
108  auto result1 = machine.transitionFromTo<Running, Completed>(42);
109  KJ_EXPECT(result1 == kj::none);
110  KJ_EXPECT(machine.is<Idle>());  // Still in Idle
111 
112  // Transition from correct state succeeds
113  machine.transitionTo<Running>(kj::str("task"));
114  auto result2 = machine.transitionFromTo<Running, Completed>(100);
115  KJ_EXPECT(result2 != kj::none);
116  KJ_EXPECT(machine.is<Completed>());
117  KJ_EXPECT(machine.getUnsafe<Completed>().result == 100);
118}
119 
120KJ_TEST("StateMachine: factory create") {
121  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Running>(kj::str("task"));
122  KJ_EXPECT(machine.is<Running>());
123  KJ_EXPECT(machine.getUnsafe<Running>().taskName == "task");
124}
125 
126// Tests for uninitialized state behavior have been removed since the default
127// constructor is now private and state machines must be created via create<>().
128 
129KJ_TEST("StateMachine: works with KJ_SWITCH_ONEOF") {
130  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Running>(kj::str("test"));
131 
132  kj::String result;
133  KJ_SWITCH_ONEOF(machine) {
134    KJ_CASE_ONEOF(idle, Idle) {
135      result = kj::str("idle");
136    }
137    KJ_CASE_ONEOF(running, Running) {
138      result = kj::str("running: ", running.taskName);
139    }
140    KJ_CASE_ONEOF(completed, Completed) {
141      result = kj::str("completed: ", completed.result);
142    }
143    KJ_CASE_ONEOF(failed, Failed) {
144      result = kj::str("failed: ", failed.error);
145    }
146  }
147 
148  KJ_EXPECT(result == "running: test");
149}
150 
151KJ_TEST("StateMachine: currentStateName introspection") {
152  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Idle>();
153 
154  // Each state
155  KJ_EXPECT(machine.currentStateName() == "idle"_kj);
156 
157  machine.transitionTo<Running>(kj::str("task"));
158  KJ_EXPECT(machine.currentStateName() == "running"_kj);
159 
160  machine.transitionTo<Completed>(42);
161  KJ_EXPECT(machine.currentStateName() == "completed"_kj);
162 
163  machine.transitionTo<Failed>(kj::str("error"));
164  KJ_EXPECT(machine.currentStateName() == "failed"_kj);
165}
166 
167// =============================================================================
168// Memory Safety Tests
169// =============================================================================
170 
171KJ_TEST("StateMachine: whenState provides safe scoped access") {
172  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Running>(kj::str("task"));
173 
174  // whenState returns result and locks transitions
175  auto result = machine.whenState<Running>([](Running& r) { return r.taskName.size(); });
176  KJ_EXPECT(result != kj::none);
177  KJ_EXPECT(KJ_ASSERT_NONNULL(result) == 4);
178 
179  // Returns none for wrong state
180  auto result2 = machine.whenState<Idle>([](Idle& i) { return i.initialized; });
181  KJ_EXPECT(result2 == kj::none);
182}
183 
184KJ_TEST("StateMachine: whenState blocks transitions during callback") {
185  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Running>(kj::str("task"));
186 
187  // Cannot transition while locked
188  auto tryTransitionInCallback = [&]() {
189    machine.whenState<Running>([&](Running&) {
190      // Attempting to transition while locked should throw
191      machine.transitionTo<Completed>(42);
192    });
193  };
194  KJ_EXPECT_THROW_MESSAGE("transitions are locked", tryTransitionInCallback());
195 
196  // State should still be Running (transition was blocked)
197  KJ_EXPECT(machine.is<Running>());
198}
199 
200KJ_TEST("StateMachine: transition lock count is tracked") {
201  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Idle>();
202 
203  KJ_EXPECT(!machine.isTransitionLocked());
204 
205  {
206    auto lock1 = machine.acquireTransitionLock();
207    KJ_EXPECT(machine.isTransitionLocked());
208 
209    {
210      auto lock2 = machine.acquireTransitionLock();
211      KJ_EXPECT(machine.isTransitionLocked());
212    }
213 
214    // Still locked after inner lock released
215    KJ_EXPECT(machine.isTransitionLocked());
216  }
217 
218  // Fully unlocked
219  KJ_EXPECT(!machine.isTransitionLocked());
220}
221 
222KJ_TEST("StateMachine: void whenState returns bool") {
223  auto machine = StateMachine<Idle, Running, Completed, Failed>::create<Running>(kj::str("task"));
224 
225  bool executed = false;
226 
227  // void callback returns true when executed
228  bool result = machine.whenState<Running>([&](Running&) { executed = true; });
229  KJ_EXPECT(result == true);
230  KJ_EXPECT(executed);
231 
232  // void callback returns false when not in state
233  executed = false;
234  bool result2 = machine.whenState<Idle>([&](Idle&) { executed = true; });
235  KJ_EXPECT(result2 == false);
236  KJ_EXPECT(!executed);
237}
238 
239// =============================================================================
240// StateMachine Tests
241// =============================================================================
242 
243// Test state types for resource lifecycle tests (TerminalStates, ErrorState, ActiveState, etc.)
244struct Active {
245  static constexpr kj::StringPtr NAME KJ_UNUSED = "active"_kj;
246  kj::String resourceName;
247  explicit Active(kj::String name): resourceName(kj::mv(name)) {}
248};
249 
250struct Closed {
251  static constexpr kj::StringPtr NAME KJ_UNUSED = "closed"_kj;
252};
253 
254struct Errored {
255  static constexpr kj::StringPtr NAME KJ_UNUSED = "errored"_kj;
256  kj::String reason;
257  explicit Errored(kj::String r): reason(kj::mv(r)) {}
258};
259 
260KJ_TEST("StateMachine: basic usage without specs") {
261  auto machine = StateMachine<Active, Closed, Errored>::create<Active>(kj::str("resource"));
262 
263  // Basic state operations work
264  KJ_EXPECT(machine.isInitialized());
265  KJ_EXPECT(machine.is<Active>());
266  KJ_EXPECT(machine.getUnsafe<Active>().resourceName == "resource");
267 
268  machine.transitionTo<Closed>();
269  KJ_EXPECT(machine.is<Closed>());
270 
271  // Can transition back (no terminal enforcement without spec)
272  machine.transitionTo<Active>(kj::str("another"));
273  KJ_EXPECT(machine.is<Active>());
274}
275 
276// Tests for uninitialized state behavior have been removed since the default
277// constructor is now private and state machines must be created via create<>().
278 
279KJ_TEST("StateMachine: with TerminalStates spec") {
280  auto machine =
281      StateMachine<TerminalStates<Closed, Errored>, Active, Closed, Errored>::create<Active>(
282          kj::str("resource"));
283  KJ_EXPECT(!machine.isTerminal());
284 
285  machine.transitionTo<Closed>();
286  KJ_EXPECT(machine.isTerminal());
287 
288  // Cannot transition from terminal state
289  auto tryTransition = [&]() { machine.transitionTo<Active>(kj::str("another")); };
290  KJ_EXPECT_THROW_MESSAGE("Cannot transition from terminal state", tryTransition());
291 
292  // But forceTransitionTo works
293  machine.forceTransitionTo<Active>(kj::str("forced"));
294  KJ_EXPECT(machine.is<Active>());
295}
296 
297KJ_TEST("StateMachine: with ErrorState spec") {
298  auto machine = StateMachine<ErrorState<Errored>, Active, Closed, Errored>::create<Active>(
299      kj::str("resource"));
300  KJ_EXPECT(!machine.isErrored());
301  KJ_EXPECT(machine.tryGetErrorUnsafe() == kj::none);
302 
303  machine.transitionTo<Errored>(kj::str("something went wrong"));
304  KJ_EXPECT(machine.isErrored());
305 
306  KJ_IF_SOME(err, machine.tryGetErrorUnsafe()) {
307    KJ_EXPECT(err.reason == "something went wrong");
308  } else {
309    KJ_FAIL_EXPECT("Should have gotten error");
310  }
311 
312  KJ_EXPECT(machine.getErrorUnsafe().reason == "something went wrong");
313}
314 
315KJ_TEST("StateMachine: with ActiveState spec") {
316  auto machine = StateMachine<ActiveState<Active>, Active, Closed, Errored>::create<Active>(
317      kj::str("resource"));
318  KJ_EXPECT(machine.isActive());
319  KJ_EXPECT(!machine.isInactive());
320 
321  KJ_IF_SOME(active, machine.tryGetActiveUnsafe()) {
322    KJ_EXPECT(active.resourceName == "resource");
323  } else {
324    KJ_FAIL_EXPECT("Should be active");
325  }
326 
327  // whenActive executes and returns value
328  auto result = machine.whenActive([](Active& a) { return a.resourceName.size(); });
329  KJ_EXPECT(result != kj::none);
330  KJ_EXPECT(KJ_ASSERT_NONNULL(result) == 8);  // "resource"
331 
332  machine.transitionTo<Closed>();
333  KJ_EXPECT(!machine.isActive());
334  KJ_EXPECT(machine.isInactive());
335 
336  // whenActive returns none when not active
337  auto result2 = machine.whenActive([](Active& a) { return a.resourceName.size(); });
338  KJ_EXPECT(result2 == kj::none);
339}
340 
341KJ_TEST("StateMachine: whenActiveOr") {
342  auto machine = StateMachine<ActiveState<Active>, Active, Closed, Errored>::create<Active>(
343      kj::str("resource"));
344 
345  // whenActiveOr executes when active
346  auto result = machine.whenActiveOr([](Active& a) { return a.resourceName.size(); }, 0ul);
347  KJ_EXPECT(result == 8);
348 
349  // After close, returns default
350  machine.transitionTo<Closed>();
351  auto result2 = machine.whenActiveOr([](Active& a) { return a.resourceName.size(); }, 999ul);
352  KJ_EXPECT(result2 == 999);
353}
354 
355KJ_TEST("StateMachine: requireActiveUnsafe") {
356  auto machine = StateMachine<ActiveState<Active>, Active, Closed, Errored>::create<Active>(
357      kj::str("resource"));
358 
359  // requireActiveUnsafeUnsafe returns reference when active
360  auto& active = machine.requireActiveUnsafe();
361  KJ_EXPECT(active.resourceName == "resource");
362 
363  // requireActiveUnsafe with custom message works when active
364  auto& active2 = machine.requireActiveUnsafe("Custom message");
365  KJ_EXPECT(active2.resourceName == "resource");
366 
367  machine.transitionTo<Closed>();
368 
369  // requireActiveUnsafe throws when not active
370  KJ_EXPECT_THROW_MESSAGE(
371      "State machine is not in the active state", (void)machine.requireActiveUnsafe());
372 
373  // requireActiveUnsafe throws custom message when not active
374  KJ_EXPECT_THROW_MESSAGE(
375      "Stream is closed", (void)machine.requireActiveUnsafe("Stream is closed"));
376}
377 
378KJ_TEST("StateMachine: with PendingStates spec") {
379  auto machine =
380      StateMachine<PendingStates<Closed, Errored>, Active, Closed, Errored>::create<Active>(
381          kj::str("resource"));
382 
383  // Start an operation
384  machine.beginOperation();
385  KJ_EXPECT(machine.hasOperationInProgress());
386 
387  // Defer a close
388  bool immediate = machine.deferTransitionTo<Closed>();
389  KJ_EXPECT(!immediate);            // Deferred
390  KJ_EXPECT(machine.is<Active>());  // Still active
391  KJ_EXPECT(machine.hasPendingState());
392  KJ_EXPECT(machine.pendingStateIs<Closed>());
393  KJ_EXPECT(machine.isOrPending<Closed>());
394 
395  // End operation - pending state applied
396  bool applied = machine.endOperation();
397  KJ_EXPECT(applied);
398  KJ_EXPECT(machine.is<Closed>());
399  KJ_EXPECT(!machine.hasPendingState());
400}
401 
402KJ_TEST("StateMachine: with PendingStates scoped operation") {
403  auto machine =
404      StateMachine<PendingStates<Closed, Errored>, Active, Closed, Errored>::create<Active>(
405          kj::str("resource"));
406 
407  {
408    auto scope = machine.scopedOperation();
409    KJ_EXPECT(machine.hasOperationInProgress());
410 
411    auto _ KJ_UNUSED = machine.deferTransitionTo<Closed>();
412    KJ_EXPECT(machine.is<Active>());  // Still active in scope
413  }
414 
415  // Scope ended, pending state applied
416  KJ_EXPECT(machine.is<Closed>());
417}
418 
419KJ_TEST("StateMachine: full-featured stream-like usage") {
420  // This demonstrates the common stream pattern with all features
421  auto machine = StateMachine<TerminalStates<Closed, Errored>, ErrorState<Errored>,
422      ActiveState<Active>, PendingStates<Closed, Errored>, Active, Closed,
423      Errored>::create<Active>(kj::str("http-body"));
424  KJ_EXPECT(machine.isActive());
425  KJ_EXPECT(!machine.isTerminal());
426  KJ_EXPECT(!machine.isErrored());
427 
428  // Safe access with whenActive
429  machine.whenActive([](Active& a) { a.resourceName = kj::str("modified"); });
430  KJ_EXPECT(machine.getUnsafe<Active>().resourceName == "modified");
431 
432  // Start a read operation
433  machine.beginOperation();
434 
435  // Close is requested mid-operation - deferred
436  auto deferred KJ_UNUSED = machine.deferTransitionTo<Closed>();
437  KJ_EXPECT(machine.isActive());  // Still active!
438  KJ_EXPECT(machine.isOrPending<Closed>());
439  KJ_EXPECT(!machine.isTerminal());  // Not terminal yet
440 
441  // End operation - close applied
442  auto applied KJ_UNUSED = machine.endOperation();
443  KJ_EXPECT(machine.is<Closed>());
444  KJ_EXPECT(machine.isTerminal());
445  KJ_EXPECT(!machine.isActive());
446  KJ_EXPECT(machine.isInactive());
447 
448  // Cannot transition from terminal
449  auto tryTransition = [&]() { machine.transitionTo<Active>(kj::str("x")); };
450  KJ_EXPECT_THROW_MESSAGE("Cannot transition from terminal state", tryTransition());
451}
452 
453KJ_TEST("StateMachine: KJ_SWITCH_ONEOF works") {
454  auto machine = StateMachine<Active, Closed, Errored>::create<Active>(kj::str("test"));
455 
456  kj::String result;
457  KJ_SWITCH_ONEOF(machine) {
458    KJ_CASE_ONEOF(active, Active) {
459      result = kj::str("active: ", active.resourceName);
460    }
461    KJ_CASE_ONEOF(closed, Closed) {
462      result = kj::str("closed");
463    }
464    KJ_CASE_ONEOF(errored, Errored) {
465      result = kj::str("errored: ", errored.reason);
466    }
467  }
468 
469  KJ_EXPECT(result == "active: test");
470}
471 
472KJ_TEST("StateMachine: whenState locks transitions") {
473  auto machine = StateMachine<Active, Closed, Errored>::create<Active>(kj::str("resource"));
474 
475  // Cannot transition while locked
476  auto tryTransitionInCallback = [&]() {
477    machine.whenState<Active>([&](Active&) { machine.transitionTo<Closed>(); });
478  };
479  KJ_EXPECT_THROW_MESSAGE("transitions are locked", tryTransitionInCallback());
480 
481  // State unchanged
482  KJ_EXPECT(machine.is<Active>());
483}
484 
485KJ_TEST("StateMachine: currentStateName") {
486  auto machine = StateMachine<Active, Closed, Errored>::create<Active>(kj::str("x"));
487  KJ_EXPECT(machine.currentStateName() == "active"_kj);
488 
489  machine.transitionTo<Closed>();
490  KJ_EXPECT(machine.currentStateName() == "closed"_kj);
491 
492  machine.transitionTo<Errored>(kj::str("err"));
493  KJ_EXPECT(machine.currentStateName() == "errored"_kj);
494}
495 
496KJ_TEST("StateMachine: const whenState works") {
497  auto machine = StateMachine<Active, Closed, Errored>::create<Active>(kj::str("resource"));
498 
499  const auto& constMachine = machine;
500 
501  // Const whenState works and returns value
502  auto result =
503      constMachine.whenState<Active>([](const Active& a) { return a.resourceName.size(); });
504  KJ_EXPECT(result != kj::none);
505  KJ_EXPECT(KJ_ASSERT_NONNULL(result) == 8);  // "resource"
506 
507  // Const whenState returns none for wrong state
508  auto result2 = constMachine.whenState<Closed>([](const Closed&) { return 42; });
509  KJ_EXPECT(result2 == kj::none);
510}
511 
512KJ_TEST("StateMachine: deferTransitionTo respects terminal states") {
513  auto machine = StateMachine<TerminalStates<Closed, Errored>, PendingStates<Closed, Errored>,
514      Active, Closed, Errored>::create<Active>(kj::str("resource"));
515 
516  // Close the machine (terminal state)
517  machine.transitionTo<Closed>();
518  KJ_EXPECT(machine.isTerminal());
519 
520  // deferTransitionTo should also fail from terminal state
521  auto tryDeferTransition = [&]() {
522    auto _ KJ_UNUSED = machine.deferTransitionTo<Errored>(kj::str("error"));
523  };
524  KJ_EXPECT_THROW_MESSAGE("Cannot transition from terminal state", tryDeferTransition());
525}
526 
527// =============================================================================
528// Streams Integration Example
529// =============================================================================
530// This demonstrates how StateMachine could replace the separate
531// state + readState pattern found in ReadableStreamInternalController.
532 
533namespace stream_integration_example {
534 
535// Simulated stream source (like ReadableStreamSource)
536struct MockSource {
537  bool dataAvailable = true;
538 
539  kj::Maybe<kj::String> read() {
540    if (dataAvailable) {
541      dataAvailable = false;
542      return kj::str("data chunk");
543    }
544    return kj::none;
545  }
546};
547 
548// State types matching the streams pattern
549struct Readable {
550  static constexpr kj::StringPtr NAME KJ_UNUSED = "readable"_kj;
551  kj::Own<MockSource> source;
552 
553  explicit Readable(kj::Own<MockSource> s): source(kj::mv(s)) {}
554};
555 
556struct StreamClosed {
557  static constexpr kj::StringPtr NAME KJ_UNUSED = "closed"_kj;
558};
559 
560struct StreamErrored {
561  static constexpr kj::StringPtr NAME KJ_UNUSED = "errored"_kj;
562  kj::String reason;
563 
564  explicit StreamErrored(kj::String r): reason(kj::mv(r)) {}
565};
566 
567// Lock states (separate state machine in the real code)
568struct Unlocked {
569  static constexpr kj::StringPtr NAME KJ_UNUSED = "unlocked"_kj;
570};
571 
572struct Locked {
573  static constexpr kj::StringPtr NAME KJ_UNUSED = "locked"_kj;
574};
575 
576struct ReaderLocked {
577  static constexpr kj::StringPtr NAME KJ_UNUSED = "reader_locked"_kj;
578  uint32_t readerId;
579  explicit ReaderLocked(uint32_t id): readerId(id) {}
580};
581 
582// The full-featured state machine type for stream data state
583using StreamDataState = StateMachine<TerminalStates<StreamClosed, StreamErrored>,
584    ErrorState<StreamErrored>,
585    ActiveState<Readable>,
586    PendingStates<StreamClosed, StreamErrored>,
587    Readable,
588    StreamClosed,
589    StreamErrored>;
590 
591// Lock state machine (simpler)
592using StreamLockState = StateMachine<Unlocked, Locked, ReaderLocked>;
593 
594// Simulated controller showing combined usage
595class MockReadableStreamController {
596 public:
597  MockReadableStreamController()
598      : dataState(StreamDataState::create<Readable>(kj::heap<MockSource>())),
599        lockState(StreamLockState::create<Unlocked>()) {}
600 
601  explicit MockReadableStreamController(kj::Own<MockSource> source)
602      : dataState(StreamDataState::create<Readable>(kj::mv(source))),
603        lockState(StreamLockState::create<Unlocked>()) {}
604 
605  bool isReadable() const {
606    return dataState.isActive();
607  }
608 
609  bool isClosedOrErrored() const {
610    return dataState.isTerminal();
611  }
612 
613  bool isErrored() const {
614    return dataState.isErrored();
615  }
616 
617  bool isLocked() const {
618    return !lockState.is<Unlocked>();
619  }
620 
621  kj::Maybe<kj::String> read() {
622    // Only read if in readable state and not already reading
623    if (!dataState.isActive()) {
624      return kj::none;
625    }
626 
627    // Start read operation (defers close/error during read)
628    auto op = dataState.scopedOperation();
629 
630    // Safe access to source
631    KJ_IF_SOME(result, dataState.whenActive([](Readable& r) -> kj::Maybe<kj::String> {
632      return r.source->read();
633    })) {
634      return kj::mv(result);
635    }
636    return kj::none;
637  }
638 
639  void close() {
640    if (dataState.isTerminal()) return;
641 
642    // If operation in progress, defer the close
643    auto _ KJ_UNUSED = dataState.deferTransitionTo<StreamClosed>();
644  }
645 
646  void error(kj::String reason) {
647    if (dataState.isTerminal()) return;
648 
649    // Error takes precedence - force even if operation in progress
650    dataState.forceTransitionTo<StreamErrored>(kj::mv(reason));
651  }
652 
653  bool acquireReaderLock(uint32_t readerId) {
654    if (isLocked()) return false;
655    lockState.transitionTo<ReaderLocked>(readerId);
656    return true;
657  }
658 
659  void releaseReaderLock() {
660    lockState.transitionTo<Unlocked>();
661  }
662 
663 private:
664  StreamDataState dataState;
665  StreamLockState lockState;
666};
667 
668}  // namespace stream_integration_example
669 
670KJ_TEST("StateMachine: stream integration example - basic flow") {
671  using namespace stream_integration_example;
672 
673  MockReadableStreamController controller(kj::heap<MockSource>());
674 
675  KJ_EXPECT(controller.isReadable());
676  KJ_EXPECT(!controller.isClosedOrErrored());
677  KJ_EXPECT(!controller.isLocked());
678 
679  // Acquire reader lock
680  KJ_EXPECT(controller.acquireReaderLock(123));
681  KJ_EXPECT(controller.isLocked());
682 
683  // Read data
684  auto chunk1 = controller.read();
685  KJ_EXPECT(chunk1 != kj::none);
686  KJ_EXPECT(KJ_ASSERT_NONNULL(chunk1) == "data chunk");
687 
688  // Second read returns none (source exhausted)
689  auto chunk2 = controller.read();
690  KJ_EXPECT(chunk2 == kj::none);
691 
692  // Close the stream
693  controller.close();
694  KJ_EXPECT(!controller.isReadable());
695  KJ_EXPECT(controller.isClosedOrErrored());
696 
697  // Release lock
698  controller.releaseReaderLock();
699  KJ_EXPECT(!controller.isLocked());
700}
701 
702KJ_TEST("StateMachine: stream integration example - close during read") {
703  using namespace stream_integration_example;
704 
705  MockReadableStreamController controller(kj::heap<MockSource>());
706 
707  // This test demonstrates that if close() is called during a read operation,
708  // the close is deferred until the read completes.
709  //
710  // In a real implementation, this would be more complex with async operations,
711  // but the pattern is the same.
712 
713  // Simulate close being called while readable (no operation in progress)
714  controller.close();
715  KJ_EXPECT(controller.isClosedOrErrored());
716}
717 
718KJ_TEST("StateMachine: stream integration example - error handling") {
719  using namespace stream_integration_example;
720 
721  MockReadableStreamController controller(kj::heap<MockSource>());
722 
723  // Error the stream
724  controller.error(kj::str("Network failure"));
725 
726  KJ_EXPECT(!controller.isReadable());
727  KJ_EXPECT(controller.isClosedOrErrored());
728  KJ_EXPECT(controller.isErrored());
729 
730  // Reads after error return none
731  auto chunk = controller.read();
732  KJ_EXPECT(chunk == kj::none);
733}
734 
735// =============================================================================
736// StateMachine Additional API Tests
737// =============================================================================
738 
739KJ_TEST("StateMachine: visit method") {
740  auto machine = StateMachine<Active, Closed, Errored>::create<Active>(kj::str("resource"));
741 
742  // Visit with return value - note: visitor must return the same type for all states
743  size_t result = machine.visit([](auto& s) -> size_t {
744    using S = std::decay_t<decltype(s)>;
745    if constexpr (std::is_same_v<S, Active>) {
746      return s.resourceName.size();
747    } else if constexpr (std::is_same_v<S, Closed>) {
748      return 0;
749    } else {
750      return s.reason.size();
751    }
752  });
753  KJ_EXPECT(result == 8);  // "resource"
754 
755  machine.transitionTo<Closed>();
756  result = machine.visit([](auto& s) -> size_t {
757    using S = std::decay_t<decltype(s)>;
758    if constexpr (std::is_same_v<S, Active>) {
759      return s.resourceName.size();
760    } else if constexpr (std::is_same_v<S, Closed>) {
761      return 0;
762    } else {
763      return s.reason.size();
764    }
765  });
766  KJ_EXPECT(result == 0);
767}
768 
769KJ_TEST("StateMachine: visit const method") {
770  auto machine = StateMachine<Active, Closed, Errored>::create<Active>(kj::str("test"));
771 
772  const auto& constMachine = machine;
773  size_t result = constMachine.visit([](const auto& s) -> size_t {
774    using S = std::decay_t<decltype(s)>;
775    if constexpr (std::is_same_v<S, Active>) {
776      return 1;
777    } else if constexpr (std::is_same_v<S, Closed>) {
778      return 2;
779    } else {
780      return 3;
781    }
782  });
783  KJ_EXPECT(result == 1);
784}
785 
786KJ_TEST("StateMachine: underlying accessor") {
787  auto machine = StateMachine<Active, Closed, Errored>::create<Active>(kj::str("resource"));
788 
789  // Access underlying kj::OneOf
790  auto& underlying = machine.underlying();
791  KJ_EXPECT(underlying.is<Active>());
792  KJ_EXPECT(underlying.get<Active>().resourceName == "resource"_kj);
793 
794  // Const access
795  const auto& constMachine = machine;
796  const auto& constUnderlying = constMachine.underlying();
797  KJ_EXPECT(constUnderlying.is<Active>());
798}
799 
800KJ_TEST("StateMachine: applyPendingStateImpl respects terminal") {
801  // When we force-transition to a terminal state during an operation,
802  // the pending state should be discarded on endOperation.
803  auto machine = StateMachine<TerminalStates<Closed, Errored>, PendingStates<Closed, Errored>,
804      Active, Closed, Errored>::create<Active>(kj::str("resource"));
805 
806  // Start an operation
807  machine.beginOperation();
808 
809  // Request a deferred close
810  auto _ KJ_UNUSED = machine.deferTransitionTo<Closed>();
811  KJ_EXPECT(machine.hasPendingState());
812  KJ_EXPECT(machine.is<Active>());
813 
814  // Force transition to error (terminal state) while operation is in progress
815  machine.forceTransitionTo<Errored>(kj::str("forced error"));
816  KJ_EXPECT(machine.is<Errored>());
817 
818  // End operation - pending Close should be discarded since we're in terminal state
819  bool pendingApplied = machine.endOperation();
820  KJ_EXPECT(!pendingApplied);             // Pending was discarded, not applied
821  KJ_EXPECT(machine.is<Errored>());       // Still in errored state
822  KJ_EXPECT(!machine.hasPendingState());  // Pending was cleared
823}
824 
825KJ_TEST("StateMachine: endOperation inside whenState throws") {
826  // This test verifies that ending an operation (which could apply a pending state)
827  // inside a whenState() callback throws an error. This prevents UAF where a
828  // transition invalidates the reference being used in the callback.
829  auto machine =
830      StateMachine<PendingStates<Closed, Errored>, Active, Closed, Errored>::create<Active>(
831          kj::str("resource"));
832 
833  // This pattern would cause UAF without the safety check:
834  //   whenState gets reference to Active
835  //   scopedOperation ends, applies pending state -> Active is destroyed
836  //   callback continues using destroyed Active reference
837  auto tryUnsafePattern = [&]() {
838    machine.whenState<Active>([&](Active&) {
839      {
840        auto op = machine.scopedOperation();
841        auto _ KJ_UNUSED = machine.deferTransitionTo<Closed>();
842      }  // op destroyed here - endOperation() would apply pending state
843    });
844  };
845 
846  KJ_EXPECT_THROW_MESSAGE("transitions are locked", tryUnsafePattern());
847 
848  // Verify the machine is still in a valid state (transition was blocked)
849  KJ_EXPECT(machine.is<Active>());
850}
851 
852KJ_TEST("StateMachine: endOperation outside whenState works") {
853  // Verify the correct pattern still works: end operations outside whenState
854  auto machine =
855      StateMachine<PendingStates<Closed, Errored>, Active, Closed, Errored>::create<Active>(
856          kj::str("resource"));
857 
858  {
859    auto op = machine.scopedOperation();
860    machine.whenState<Active>([&](Active& a) {
861      // Safe to use 'a' here - no operation ending in this scope
862      KJ_EXPECT(a.resourceName == "resource");
863    });
864    auto _ KJ_UNUSED = machine.deferTransitionTo<Closed>();
865  }  // op ends here, OUTSIDE any whenState callback - safe!
866 
867  KJ_EXPECT(machine.is<Closed>());
868}
869 
870}  // namespace
871}  // namespace workerd