diff --git a/src/workerd/api/container-test.c++ b/src/workerd/api/container-test.c++ index e47f5e9a16a..42b1dfe668b 100644 --- a/src/workerd/api/container-test.c++ +++ b/src/workerd/api/container-test.c++ @@ -231,6 +231,47 @@ class MockExecContainerServer final: public rpc::Container::Server { kj::Maybe>> resizeFulfiller; }; +struct CapturedDirectorySnapshot { + bool hasSnapshotId = false; + kj::String snapshotId; + kj::String restorePath; +}; + +// Captures the directorySnapshots forwarded by Container::start() so tests can assert how each +// DirectorySnapshotRestoreParams is translated into the RPC request. +class DirectorySnapshotStartServer final: public rpc::Container::Server { + public: + explicit DirectorySnapshotStartServer(kj::Vector& captured) + : captured(captured) {} + + kj::Promise start(StartContext context) override { + for (auto entry: context.getParams().getDirectorySnapshots()) { + captured.add(CapturedDirectorySnapshot{ + .hasSnapshotId = entry.hasSnapshotId(), + .snapshotId = kj::str(entry.getSnapshotId()), + .restorePath = kj::str(entry.getRestorePath()), + }); + } + return kj::READY_NOW; + } + + kj::Promise monitor(MonitorContext context) override { + return kj::NEVER_DONE; + } + + private: + kj::Vector& captured; +}; + +Container::DirectorySnapshot makeDirectorySnapshot(kj::StringPtr id, kj::StringPtr dir) { + return Container::DirectorySnapshot{ + .id = kj::str(id), + .size = 0, + .dir = kj::str(dir), + .name = kj::none, + }; +} + enum class TunnelReuseGate { DISABLED, ENABLED }; class AutogateScope { @@ -748,6 +789,126 @@ KJ_TEST("Container::snapshotContainer propagates the current span context") { KJ_EXPECT(containerCalled); } +KJ_TEST("Container::start restores a directory snapshot using the snapshot's own dir") { + kj::Vector captured; + auto fixture = makeFixture(); + + fixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise { + auto container = kj::heap( + rpc::Container::Client(kj::heap(captured)), false); + + auto snapshots = kj::heapArrayBuilder(1); + snapshots.add(Container::DirectorySnapshotRestoreParams{ + .snapshot = makeDirectorySnapshot("snap-id"_kj, "/data"_kj), + .mountPoint = kj::none, + }); + container->start(env.js, + Container::StartupOptions{ + .directorySnapshots = snapshots.finish(), + }); + + // Give the queued start RPC bounded time to run. + for (auto i = 0; i < 10; ++i) { + co_await kj::yield(); + } + }); + + KJ_ASSERT(captured.size() == 1); + KJ_EXPECT(captured[0].hasSnapshotId); + KJ_EXPECT(captured[0].snapshotId == "snap-id"); + KJ_EXPECT(captured[0].restorePath == "/data"); +} + +KJ_TEST("Container::start lets mountPoint override the snapshot's dir as the restore path") { + kj::Vector captured; + auto fixture = makeFixture(); + + fixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise { + auto container = kj::heap( + rpc::Container::Client(kj::heap(captured)), false); + + auto snapshots = kj::heapArrayBuilder(1); + snapshots.add(Container::DirectorySnapshotRestoreParams{ + .snapshot = makeDirectorySnapshot("snap-id"_kj, "/data"_kj), + .mountPoint = kj::str("/mnt/elsewhere"), + }); + container->start(env.js, + Container::StartupOptions{ + .directorySnapshots = snapshots.finish(), + }); + + for (auto i = 0; i < 10; ++i) { + co_await kj::evalLater([]() {}); + } + }); + + KJ_ASSERT(captured.size() == 1); + KJ_EXPECT(captured[0].hasSnapshotId); + KJ_EXPECT(captured[0].snapshotId == "snap-id"); + KJ_EXPECT(captured[0].restorePath == "/mnt/elsewhere"); +} + +KJ_TEST("Container::start restores to mountPoint when no snapshot is given") { + kj::Vector captured; + auto fixture = makeFixture(); + + fixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise { + auto container = kj::heap( + rpc::Container::Client(kj::heap(captured)), false); + + auto snapshots = kj::heapArrayBuilder(1); + snapshots.add(Container::DirectorySnapshotRestoreParams{ + .snapshot = kj::none, + .mountPoint = kj::str("/mnt/data"), + }); + container->start(env.js, + Container::StartupOptions{ + .directorySnapshots = snapshots.finish(), + }); + + for (auto i = 0; i < 10; ++i) { + co_await kj::evalLater([]() {}); + } + }); + + KJ_ASSERT(captured.size() == 1); + KJ_EXPECT(!captured[0].hasSnapshotId); + KJ_EXPECT(captured[0].snapshotId == ""); + KJ_EXPECT(captured[0].restorePath == "/mnt/data"); +} + +KJ_TEST("Container::start requires mountPoint when no snapshot is given") { + kj::Vector captured; + auto fixture = makeFixture(); + + fixture.runInIoContext([&](const TestFixture::Environment& env) -> kj::Promise { + auto container = kj::heap( + rpc::Container::Client(kj::heap(captured)), false); + + auto snapshots = kj::heapArrayBuilder(1); + snapshots.add(Container::DirectorySnapshotRestoreParams{ + .snapshot = kj::none, + .mountPoint = kj::none, + }); + + bool threw = false; + JSG_TRY(env.js) { + container->start(env.js, + Container::StartupOptions{ + .directorySnapshots = snapshots.finish(), + }); + } + JSG_CATCH(exception KJ_UNUSED) { + threw = true; + } + KJ_EXPECT(threw, "start() should throw when neither snapshot nor mountPoint is given"); + + return kj::READY_NOW; + }); + + KJ_EXPECT(captured.size() == 0); +} + KJ_TEST("Container::exec forwards pty options and resize() sends a resize RPC") { ExecObservations observations; auto fixture = makeFixture(); diff --git a/src/workerd/api/container.c++ b/src/workerd/api/container.c++ index 71f3389e2b0..51a5c906b07 100644 --- a/src/workerd/api/container.c++ +++ b/src/workerd/api/container.c++ @@ -420,17 +420,25 @@ void Container::start(jsg::Lock& js, jsg::Optional maybeOptions) for (auto i: kj::indices(directorySnapshots)) { auto entry = list[i]; auto& restore = directorySnapshots[i]; - auto& snap = restore.snapshot; - auto effectiveRestorePath = snap.dir.asPtr(); + + kj::Maybe effectiveRestorePath; + KJ_IF_SOME(snap, restore.snapshot) { + effectiveRestorePath = snap.dir.asPtr(); + } KJ_IF_SOME(mp, restore.mountPoint) { effectiveRestorePath = mp.asPtr(); } - JSG_REQUIRE_NONNULL(parseRestorePath(effectiveRestorePath), Error, + auto restorePath = JSG_REQUIRE_NONNULL(effectiveRestorePath, Error, + "Directory snapshot restore requires a mountPoint when no snapshot is given."); + + JSG_REQUIRE_NONNULL(parseRestorePath(restorePath), Error, "Directory snapshot cannot be restored to root directory."); - entry.setSnapshotId(snap.id); - entry.setRestorePath(effectiveRestorePath); + KJ_IF_SOME(snap, restore.snapshot) { + entry.setSnapshotId(snap.id); + } + entry.setRestorePath(restorePath); } } diff --git a/src/workerd/api/container.h b/src/workerd/api/container.h index d097b597211..c7d03a4cb45 100644 --- a/src/workerd/api/container.h +++ b/src/workerd/api/container.h @@ -215,10 +215,14 @@ class Container: public jsg::Object { }; struct DirectorySnapshotRestoreParams { - DirectorySnapshot snapshot; + jsg::Optional snapshot; jsg::Optional mountPoint; JSG_STRUCT(snapshot, mountPoint); + JSG_STRUCT_TS_OVERRIDE(type ContainerDirectorySnapshotRestoreParams = + | { snapshot: ContainerDirectorySnapshot; mountPoint?: string } + | { snapshot?: undefined; mountPoint: string } + ); }; struct Snapshot { diff --git a/types/generated-snapshot/experimental/index.d.ts b/types/generated-snapshot/experimental/index.d.ts index 7c1db03ea1d..01719ef4c8b 100755 --- a/types/generated-snapshot/experimental/index.d.ts +++ b/types/generated-snapshot/experimental/index.d.ts @@ -4081,10 +4081,15 @@ interface ContainerDirectorySnapshotOptions { dir: string; name?: string; } -interface ContainerDirectorySnapshotRestoreParams { - snapshot: ContainerDirectorySnapshot; - mountPoint?: string; -} +type ContainerDirectorySnapshotRestoreParams = + | { + snapshot: ContainerDirectorySnapshot; + mountPoint?: string; + } + | { + snapshot?: undefined; + mountPoint: string; + }; interface ContainerSnapshot { id: string; size: number; diff --git a/types/generated-snapshot/experimental/index.ts b/types/generated-snapshot/experimental/index.ts index 19c5a115626..51da200634a 100755 --- a/types/generated-snapshot/experimental/index.ts +++ b/types/generated-snapshot/experimental/index.ts @@ -4090,10 +4090,15 @@ export interface ContainerDirectorySnapshotOptions { dir: string; name?: string; } -export interface ContainerDirectorySnapshotRestoreParams { - snapshot: ContainerDirectorySnapshot; - mountPoint?: string; -} +export type ContainerDirectorySnapshotRestoreParams = + | { + snapshot: ContainerDirectorySnapshot; + mountPoint?: string; + } + | { + snapshot?: undefined; + mountPoint: string; + }; export interface ContainerSnapshot { id: string; size: number; diff --git a/types/generated-snapshot/index.d.ts b/types/generated-snapshot/index.d.ts index 86619d47828..f6f42dd3713 100755 --- a/types/generated-snapshot/index.d.ts +++ b/types/generated-snapshot/index.d.ts @@ -3991,10 +3991,15 @@ interface ContainerDirectorySnapshotOptions { dir: string; name?: string; } -interface ContainerDirectorySnapshotRestoreParams { - snapshot: ContainerDirectorySnapshot; - mountPoint?: string; -} +type ContainerDirectorySnapshotRestoreParams = + | { + snapshot: ContainerDirectorySnapshot; + mountPoint?: string; + } + | { + snapshot?: undefined; + mountPoint: string; + }; interface ContainerSnapshot { id: string; size: number; diff --git a/types/generated-snapshot/index.ts b/types/generated-snapshot/index.ts index e61d2a0627a..99e938d922c 100755 --- a/types/generated-snapshot/index.ts +++ b/types/generated-snapshot/index.ts @@ -4000,10 +4000,15 @@ export interface ContainerDirectorySnapshotOptions { dir: string; name?: string; } -export interface ContainerDirectorySnapshotRestoreParams { - snapshot: ContainerDirectorySnapshot; - mountPoint?: string; -} +export type ContainerDirectorySnapshotRestoreParams = + | { + snapshot: ContainerDirectorySnapshot; + mountPoint?: string; + } + | { + snapshot?: undefined; + mountPoint: string; + }; export interface ContainerSnapshot { id: string; size: number;