fix(harness): defer self-managed sandbox shutdown until session mirro… #3094
There are no files selected for viewing
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -21,7 +21,9 @@ | |
| import io.agentscope.harness.agent.filesystem.model.ReadResult; | ||
| import io.agentscope.harness.agent.filesystem.sandbox.PinnedSandboxFilesystem; | ||
| import io.agentscope.harness.agent.sandbox.Sandbox; | ||
| import io.agentscope.harness.agent.sandbox.SandboxAcquireResult; | ||
| import io.agentscope.harness.agent.sandbox.SandboxAware; | ||
| import io.agentscope.harness.agent.sandbox.SandboxMirrorReleaseCoordinator; | ||
| import io.agentscope.harness.agent.transcript.ObjectStoreTranscriptStore; | ||
| import io.agentscope.harness.agent.transcript.TranscriptRef; | ||
| import io.agentscope.harness.agent.transcript.TranscriptStore; | ||
|
|
@@ -112,9 +114,15 @@ public class SessionTree { | |
| * @return {@code true} if the mirrors quiesced within the timeout | ||
| */ | ||
| public static boolean awaitMirrorQuiescence(long timeout, TimeUnit unit) { | ||
| long deadlineNanos = System.nanoTime() + unit.toNanos(timeout); | ||
| try { | ||
| long mirrorWaitNanos = Math.max(0L, deadlineNanos - System.nanoTime()); | ||
| MIRROR_EXECUTOR.submit(() -> {}).get(mirrorWaitNanos, TimeUnit.NANOSECONDS); | ||
| // Deferred stop/shutdown runs on a separate executor; drain it too so graceful close | ||
| // does not race sandbox teardown that was scheduled from mirror finally blocks. | ||
| long releaseWaitNanos = Math.max(0L, deadlineNanos - System.nanoTime()); | ||
| return SandboxMirrorReleaseCoordinator.awaitReleaseQuiescence( | ||
| releaseWaitNanos, TimeUnit.NANOSECONDS); | ||
| } catch (InterruptedException e) { | ||
| Thread.currentThread().interrupt(); | ||
| return false; | ||
|
|
@@ -550,17 +558,18 @@ private void scheduleSegmentMirror(List<SessionEntry> entries, long seqStart, lo | |
| if (transcriptStore == null || transcriptRef == null || entries.isEmpty()) { | ||
| return; | ||
| } | ||
| // Prepare work before retain so a failure here cannot leak a deferred sandbox release. | ||
| final AbstractFilesystem mirrorFs = pinIfSandbox(filesystem); | ||
| final TranscriptStore store = transcriptStoreForMirror(mirrorFs); | ||
| final TranscriptRef ref = transcriptRef; | ||
| StringBuilder sb = new StringBuilder(); | ||
| for (SessionEntry entry : entries) { | ||
| sb.append(JsonUtils.getJsonCodec().toJson(entry)).append('\n'); | ||
| } | ||
| final byte[] payload = sb.toString().getBytes(StandardCharsets.UTF_8); | ||
| final String wid = writerId; | ||
| submitPinnedMirror( | ||
| mirrorFs, | ||
| () -> { | ||
| try { | ||
| store.appendSegment(ref, seqStart, seqEnd, wid, payload); | ||
|
|
@@ -584,30 +593,90 @@ private void scheduleMirror() { | |
| if (filesystem == null || workspaceRoot == null) { | ||
| return; | ||
| } | ||
| // Resolve paths before retain so retain↔execute has no prepare-side leak window. | ||
| final AbstractFilesystem mirrorFs = pinIfSandbox(filesystem); | ||
| final String contextRel = resolveRelativePath(contextFile); | ||
| final String logRel = resolveRelativePath(logFile); | ||
| submitPinnedMirror( | ||
| mirrorFs, | ||
| () -> { | ||
| mirrorToFilesystem(mirrorFs, contextFile, contextRel); | ||
| mirrorToFilesystem(mirrorFs, logFile, logRel); | ||
| }); | ||
| } | ||
|
|
||
| /** | ||
| * Retains a pinned sandbox (if any), submits {@code task} to the mirror executor, and always | ||
| * pairs {@code releaseMirror} — either in the task {@code finally} or immediately when | ||
| * submission itself fails. | ||
| */ | ||
| private static void submitPinnedMirror(AbstractFilesystem mirrorFs, Runnable task) { | ||
| final Sandbox retainedSandbox = sandboxForMirrorRetain(mirrorFs); | ||
| if (retainedSandbox != null) { | ||
| SandboxMirrorReleaseCoordinator.retain(retainedSandbox); | ||
| } | ||
| try { | ||
| MIRROR_EXECUTOR.execute( | ||
| () -> { | ||
| try { | ||
| task.run(); | ||
| } finally { | ||
| if (retainedSandbox != null) { | ||
| SandboxMirrorReleaseCoordinator.releaseMirror(retainedSandbox); | ||
| } | ||
| } | ||
| }); | ||
| } catch (RuntimeException e) { | ||
| if (retainedSandbox != null) { | ||
| SandboxMirrorReleaseCoordinator.releaseMirror(retainedSandbox); | ||
| } | ||
| throw e; | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * When {@code fs} is a sandbox-backed filesystem, return a pinned filesystem that keeps the | ||
| * <em>per-call</em> sandbox for async uploads. Prefers {@link SandboxAcquireResult} on {@link | ||
| * #fsRc} (issue #2490) over the legacy shared {@link SandboxAware#getSandbox()} field, which | ||
| * is last-writer-wins under concurrent distinct sessions. | ||
| */ | ||
| private AbstractFilesystem pinIfSandbox(AbstractFilesystem fs) { | ||
| Sandbox sb = resolveSandboxForPin(fs); | ||
| if (sb != null) { | ||
| return new PinnedSandboxFilesystem(sb); | ||
| } | ||
| return fs; | ||
| } | ||
|
|
||
| /** | ||
| * Resolves the sandbox to pin for an async mirror: per-call {@link RuntimeContext} binding | ||
| * first, then the shared {@link SandboxAware} fallback for context-free callers. | ||
| */ | ||
| private Sandbox resolveSandboxForPin(AbstractFilesystem fs) { | ||
| RuntimeContext rc = fsRc; | ||
| if (rc != null) { | ||
| SandboxAcquireResult bound = rc.get(SandboxAcquireResult.class); | ||
| if (bound != null && bound.getSandbox() != null) { | ||
| return bound.getSandbox(); | ||
| } | ||
| } | ||
| if (fs instanceof SandboxAware aware) { | ||
| return aware.getSandbox(); | ||
| } | ||
| return null; | ||
| } | ||
|
|
||
| /** | ||
| * Returns the sandbox pinned for an async mirror task, or {@code null} when the mirror does | ||
| * not hold a sandbox connection that must defer self-managed release. | ||
| */ | ||
| private static Sandbox sandboxForMirrorRetain(AbstractFilesystem mirrorFs) { | ||
|
Collaborator
There was a problem hiding this comment.
|
||
| if (mirrorFs instanceof PinnedSandboxFilesystem pinned) { | ||
| return pinned.getSandbox(); | ||
| } | ||
| return null; | ||
| } | ||
|
|
||
| private TranscriptStore transcriptStoreForMirror(AbstractFilesystem mirrorFs) { | ||
| if (transcriptStore instanceof ObjectStoreTranscriptStore ost) { | ||
| return ost.withFilesystem(mirrorFs).withRuntimeContext(fsRc); | ||
|
|
||
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
[Info] Nice fix draining the release executor within a shared deadline. Two small notes: (1)
awaitReleaseQuiescencesubmits an empty task, so it only waits for releases enqueued before the barrier call — releases scheduled by mirror tasks that begin executing after the submit are not covered; that matches the documented "submitted before this call" contract, just make sure graceful-shutdown callers invoke it afterawaitMirrorQuiescencereturns true (which the current call ordering does). (2) The process-wide singleton coordinator means multi-agent JVMs share one release thread — fine for teardown throughput, but a slowstop()head-of-line blocks other sessions' deferred releases; acceptable given it is off the mirror path.