1 files changed, 234 insertions, 126 deletions

diff --git a/src/zencompute/runners/remotehttprunner.cpp b/src/zencompute/runners/remotehttprunner.cpp
index ce6a81173..55f78fdd6 100644
--- a/src/zencompute/runners/remotehttprunner.cpp
+++ b/src/zencompute/runners/remotehttprunner.cpp
@@ -20,6 +20,7 @@
 #	include <zenstore/cidstore.h>
 
 #	include <span>
+#	include <unordered_set>
 
 //////////////////////////////////////////////////////////////////////////
 
@@ -38,6 +39,7 @@ RemoteHttpRunner::RemoteHttpRunner(ChunkResolver&				InChunkResolver,
 , m_ChunkResolver{InChunkResolver}
 , m_WorkerPool{InWorkerPool}
 , m_HostName{HostName}
+, m_DisplayName{HostName}
 , m_BaseUrl{fmt::format("{}/compute", HostName)}
 , m_Http(m_BaseUrl)
 , m_InstanceId(Oid::NewOid())
@@ -59,6 +61,15 @@ RemoteHttpRunner::RemoteHttpRunner(ChunkResolver&				InChunkResolver,
 	m_MonitorThread = std::thread{&RemoteHttpRunner::MonitorThreadFunction, this};
 }
 
+void
+RemoteHttpRunner::SetRemoteHostname(std::string_view Hostname)
+{
+	if (!Hostname.empty())
+	{
+		m_DisplayName = fmt::format("{} ({})", m_HostName, Hostname);
+	}
+}
+
 RemoteHttpRunner::~RemoteHttpRunner()
 {
 	Shutdown();
@@ -108,6 +119,7 @@ RemoteHttpRunner::Shutdown()
 	for (auto& [RemoteLsn, HttpAction] : Remaining)
 	{
 		ZEN_DEBUG("shutdown: marking remote action LSN {} (local LSN {}) as Failed", RemoteLsn, HttpAction.Action->ActionLsn);
+		HttpAction.Action->FailureReason = "remote runner shutdown";
 		HttpAction.Action->SetActionState(RunnerAction::State::Failed);
 	}
 }
@@ -213,11 +225,13 @@ RemoteHttpRunner::QueryCapacity()
 		return 0;
 	}
 
-	// Estimate how much more work we're ready to accept
+	// Estimate how much more work we're ready to accept.
+	// Include actions currently being submitted over HTTP so we don't
+	// keep queueing new submissions while previous ones are still in flight.
 
 	RwLock::SharedLockScope _{m_RunningLock};
 
-	size_t RunningCount = m_RemoteRunningMap.size();
+	size_t RunningCount = m_RemoteRunningMap.size() + m_InFlightSubmissions.load(std::memory_order_relaxed);
 
 	if (RunningCount >= size_t(m_MaxRunningActions))
 	{
@@ -232,6 +246,9 @@ RemoteHttpRunner::SubmitActions(const std::vector<Ref<RunnerAction>>& Actions)
 {
 	ZEN_TRACE_CPU("RemoteHttpRunner::SubmitActions");
 
+	m_InFlightSubmissions.fetch_add(Actions.size(), std::memory_order_relaxed);
+	auto InFlightGuard = MakeGuard([&] { m_InFlightSubmissions.fetch_sub(Actions.size(), std::memory_order_relaxed); });
+
 	if (Actions.size() <= 1)
 	{
 		std::vector<SubmitResult> Results;
@@ -359,108 +376,141 @@ RemoteHttpRunner::SubmitAction(Ref<RunnerAction> Action)
 		}
 	}
 
-	// Enqueue job. If the remote returns FailedDependency (424), it means it
-	// cannot resolve the worker/function — re-register the worker and retry once.
+	// Submit the action to the remote. In eager-attach mode we build a
+	// CbPackage with all referenced attachments upfront to avoid the 404
+	// round-trip. In the default mode we POST the bare object first and
+	// only upload missing attachments if the remote requests them.
+	//
+	// In both modes, FailedDependency (424) triggers a worker re-register
+	// and a single retry.
 
 	CbObject			 Result;
 	HttpClient::Response WorkResponse;
 	HttpResponseCode	 WorkResponseCode{};
 
-	for (int Attempt = 0; Attempt < 2; ++Attempt)
-	{
-		WorkResponse	 = m_Http.Post(SubmitUrl, ActionObj);
-		WorkResponseCode = WorkResponse.StatusCode;
-
-		if (WorkResponseCode == HttpResponseCode::FailedDependency && Attempt == 0)
-		{
-			ZEN_WARN("remote {} returned FailedDependency for action {} — re-registering worker and retrying",
-					 m_Http.GetBaseUri(),
-					 ActionId);
-
-			(void)RegisterWorker(Action->Worker.Descriptor);
-		}
-		else
-		{
-			break;
-		}
-	}
-
-	if (WorkResponseCode == HttpResponseCode::OK)
-	{
-		Result = WorkResponse.AsObject();
-	}
-	else if (WorkResponseCode == HttpResponseCode::NotFound)
+	if (m_EagerAttach)
 	{
-		// Not all attachments are present
-
-		// Build response package including all required attachments
-
 		CbPackage Pkg;
 		Pkg.SetObject(ActionObj);
 
-		CbObject Response = WorkResponse.AsObject();
+		ActionObj.IterateAttachments([&](CbFieldView Field) {
+			const IoHash AttachHash = Field.AsHash();
 
-		for (auto& Item : Response["need"sv])
-		{
-			const IoHash NeedHash = Item.AsHash();
-
-			if (IoBuffer Chunk = m_ChunkResolver.FindChunkByCid(NeedHash))
+			if (IoBuffer Chunk = m_ChunkResolver.FindChunkByCid(AttachHash))
 			{
 				uint64_t		 DataRawSize = 0;
 				IoHash			 DataRawHash;
 				CompressedBuffer Compressed =
 					CompressedBuffer::FromCompressed(SharedBuffer{Chunk}, /* out */ DataRawHash, /* out */ DataRawSize);
 
-				ZEN_ASSERT(DataRawHash == NeedHash);
+				Pkg.AddAttachment(CbAttachment(Compressed, AttachHash));
+				m_LastSubmitStats.TotalAttachments.fetch_add(1, std::memory_order_relaxed);
+				m_LastSubmitStats.TotalAttachmentBytes.fetch_add(Chunk.GetSize(), std::memory_order_relaxed);
+			}
+		});
+
+		for (int Attempt = 0; Attempt < 2; ++Attempt)
+		{
+			WorkResponse	 = m_Http.Post(SubmitUrl, Pkg);
+			WorkResponseCode = WorkResponse.StatusCode;
+
+			if (WorkResponseCode == HttpResponseCode::FailedDependency && Attempt == 0)
+			{
+				ZEN_WARN("remote {} returned FailedDependency for action {} — re-registering worker and retrying",
+						 m_Http.GetBaseUri(),
+						 ActionId);
 
-				Pkg.AddAttachment(CbAttachment(Compressed, NeedHash));
+				(void)RegisterWorker(Action->Worker.Descriptor);
 			}
 			else
 			{
-				// No such attachment
-
-				return {.IsAccepted = false, .Reason = fmt::format("missing attachment {}", NeedHash)};
+				break;
 			}
 		}
+	}
+	else
+	{
+		for (int Attempt = 0; Attempt < 2; ++Attempt)
+		{
+			WorkResponse	 = m_Http.Post(SubmitUrl, ActionObj);
+			WorkResponseCode = WorkResponse.StatusCode;
 
-		// Post resulting package
+			if (WorkResponseCode == HttpResponseCode::FailedDependency && Attempt == 0)
+			{
+				ZEN_WARN("remote {} returned FailedDependency for action {} — re-registering worker and retrying",
+						 m_Http.GetBaseUri(),
+						 ActionId);
 
-		HttpClient::Response PayloadResponse = m_Http.Post(SubmitUrl, Pkg);
+				(void)RegisterWorker(Action->Worker.Descriptor);
+			}
+			else
+			{
+				break;
+			}
+		}
 
-		if (!PayloadResponse)
+		if (WorkResponseCode == HttpResponseCode::NotFound)
 		{
-			ZEN_WARN("unable to register payloads for action {} at {}{}", ActionId, m_Http.GetBaseUri(), SubmitUrl);
+			// Remote needs attachments — resolve them and retry with a CbPackage
 
-			// TODO: include more information about the failure in the response
+			CbPackage Pkg;
+			Pkg.SetObject(ActionObj);
 
-			return {.IsAccepted = false, .Reason = "HTTP request failed"};
-		}
-		else if (PayloadResponse.StatusCode == HttpResponseCode::OK)
-		{
-			Result = PayloadResponse.AsObject();
-		}
-		else
-		{
-			// Unexpected response
-
-			const int ResponseStatusCode = (int)PayloadResponse.StatusCode;
-
-			ZEN_WARN("unable to register payloads for action {} at {}{} (error: {} {})",
-					 ActionId,
-					 m_Http.GetBaseUri(),
-					 SubmitUrl,
-					 ResponseStatusCode,
-					 ToString(ResponseStatusCode));
-
-			return {.IsAccepted = false,
-					.Reason		= fmt::format("unexpected response code {} {} from {}{}",
-										  ResponseStatusCode,
-										  ToString(ResponseStatusCode),
-										  m_Http.GetBaseUri(),
-										  SubmitUrl)};
+			CbObject Response = WorkResponse.AsObject();
+
+			for (auto& Item : Response["need"sv])
+			{
+				const IoHash NeedHash = Item.AsHash();
+
+				if (IoBuffer Chunk = m_ChunkResolver.FindChunkByCid(NeedHash))
+				{
+					uint64_t		 DataRawSize = 0;
+					IoHash			 DataRawHash;
+					CompressedBuffer Compressed =
+						CompressedBuffer::FromCompressed(SharedBuffer{Chunk}, /* out */ DataRawHash, /* out */ DataRawSize);
+
+					ZEN_ASSERT(DataRawHash == NeedHash);
+
+					Pkg.AddAttachment(CbAttachment(Compressed, NeedHash));
+					m_LastSubmitStats.TotalAttachments.fetch_add(1, std::memory_order_relaxed);
+					m_LastSubmitStats.TotalAttachmentBytes.fetch_add(Chunk.GetSize(), std::memory_order_relaxed);
+				}
+				else
+				{
+					return {.IsAccepted = false, .Reason = fmt::format("missing attachment {}", NeedHash)};
+				}
+			}
+
+			HttpClient::Response PayloadResponse = m_Http.Post(SubmitUrl, Pkg);
+
+			if (!PayloadResponse)
+			{
+				ZEN_WARN("unable to register payloads for action {} at {}{}", ActionId, m_Http.GetBaseUri(), SubmitUrl);
+				return {.IsAccepted = false, .Reason = "HTTP request failed"};
+			}
+
+			WorkResponse	 = std::move(PayloadResponse);
+			WorkResponseCode = WorkResponse.StatusCode;
 		}
 	}
 
+	if (WorkResponseCode == HttpResponseCode::OK)
+	{
+		Result = WorkResponse.AsObject();
+	}
+	else if (!WorkResponse)
+	{
+		ZEN_WARN("submit of action {} to {}{} failed", ActionId, m_Http.GetBaseUri(), SubmitUrl);
+		return {.IsAccepted = false, .Reason = "HTTP request failed"};
+	}
+	else if (!IsHttpSuccessCode(WorkResponseCode))
+	{
+		const int Code = static_cast<int>(WorkResponseCode);
+		ZEN_WARN("submit of action {} to {}{} returned {} {}", ActionId, m_Http.GetBaseUri(), SubmitUrl, Code, ToString(Code));
+		return {.IsAccepted = false,
+				.Reason = fmt::format("unexpected response code {} {} from {}{}", Code, ToString(Code), m_Http.GetBaseUri(), SubmitUrl)};
+	}
+
 	if (Result)
 	{
 		if (const int32_t LsnField = Result["lsn"].AsInt32(0))
@@ -512,82 +562,110 @@ RemoteHttpRunner::SubmitActionBatch(const std::string& SubmitUrl, const std::vec
 	CbObjectWriter Body;
 	Body.BeginArray("actions"sv);
 
+	std::unordered_set<IoHash, IoHash::Hasher> AttachmentsSeen;
+
 	for (const Ref<RunnerAction>& Action : Actions)
 	{
 		Action->ExecutionLocation = m_HostName;
 		MaybeDumpAction(Action->ActionLsn, Action->ActionObj);
 		Body.AddObject(Action->ActionObj);
+
+		if (m_EagerAttach)
+		{
+			Action->ActionObj.IterateAttachments([&](CbFieldView Field) { AttachmentsSeen.insert(Field.AsHash()); });
+		}
 	}
 
 	Body.EndArray();
 
-	// POST the batch
-
-	HttpClient::Response Response = m_Http.Post(SubmitUrl, Body.Save());
-
-	if (Response.StatusCode == HttpResponseCode::OK)
-	{
-		return ParseBatchResponse(Response, Actions);
-	}
+	// In eager-attach mode, build a CbPackage with all referenced attachments
+	// so the remote can accept in a single round-trip.  Otherwise POST a bare
+	// CbObject and handle the 404 need-list flow.
 
-	if (Response.StatusCode == HttpResponseCode::NotFound)
+	if (m_EagerAttach)
 	{
-		// Server needs attachments — resolve them and retry with a CbPackage
-
-		CbObject NeedObj = Response.AsObject();
-
 		CbPackage Pkg;
 		Pkg.SetObject(Body.Save());
 
-		for (auto& Item : NeedObj["need"sv])
+		for (const IoHash& AttachHash : AttachmentsSeen)
 		{
-			const IoHash NeedHash = Item.AsHash();
-
-			if (IoBuffer Chunk = m_ChunkResolver.FindChunkByCid(NeedHash))
+			if (IoBuffer Chunk = m_ChunkResolver.FindChunkByCid(AttachHash))
 			{
 				uint64_t		 DataRawSize = 0;
 				IoHash			 DataRawHash;
 				CompressedBuffer Compressed =
 					CompressedBuffer::FromCompressed(SharedBuffer{Chunk}, /* out */ DataRawHash, /* out */ DataRawSize);
 
-				ZEN_ASSERT(DataRawHash == NeedHash);
-
-				Pkg.AddAttachment(CbAttachment(Compressed, NeedHash));
-			}
-			else
-			{
-				ZEN_WARN("batch submit: missing attachment {} — falling back to individual submit", NeedHash);
-				return FallbackToIndividualSubmit(Actions);
+				Pkg.AddAttachment(CbAttachment(Compressed, AttachHash));
+				m_LastSubmitStats.TotalAttachments.fetch_add(1, std::memory_order_relaxed);
+				m_LastSubmitStats.TotalAttachmentBytes.fetch_add(Chunk.GetSize(), std::memory_order_relaxed);
 			}
 		}
 
-		HttpClient::Response RetryResponse = m_Http.Post(SubmitUrl, Pkg);
+		HttpClient::Response Response = m_Http.Post(SubmitUrl, Pkg);
+
+		if (Response.StatusCode == HttpResponseCode::OK)
+		{
+			return ParseBatchResponse(Response, Actions);
+		}
+	}
+	else
+	{
+		HttpClient::Response Response = m_Http.Post(SubmitUrl, Body.Save());
 
-		if (RetryResponse.StatusCode == HttpResponseCode::OK)
+		if (Response.StatusCode == HttpResponseCode::OK)
 		{
-			return ParseBatchResponse(RetryResponse, Actions);
+			return ParseBatchResponse(Response, Actions);
 		}
 
-		ZEN_WARN("batch submit retry failed with {} {} — falling back to individual submit",
-				 (int)RetryResponse.StatusCode,
-				 ToString(RetryResponse.StatusCode));
-		return FallbackToIndividualSubmit(Actions);
+		if (Response.StatusCode == HttpResponseCode::NotFound)
+		{
+			CbObject NeedObj = Response.AsObject();
+
+			CbPackage Pkg;
+			Pkg.SetObject(Body.Save());
+
+			for (auto& Item : NeedObj["need"sv])
+			{
+				const IoHash NeedHash = Item.AsHash();
+
+				if (IoBuffer Chunk = m_ChunkResolver.FindChunkByCid(NeedHash))
+				{
+					uint64_t		 DataRawSize = 0;
+					IoHash			 DataRawHash;
+					CompressedBuffer Compressed =
+						CompressedBuffer::FromCompressed(SharedBuffer{Chunk}, /* out */ DataRawHash, /* out */ DataRawSize);
+
+					ZEN_ASSERT(DataRawHash == NeedHash);
+
+					Pkg.AddAttachment(CbAttachment(Compressed, NeedHash));
+					m_LastSubmitStats.TotalAttachments.fetch_add(1, std::memory_order_relaxed);
+					m_LastSubmitStats.TotalAttachmentBytes.fetch_add(Chunk.GetSize(), std::memory_order_relaxed);
+				}
+				else
+				{
+					ZEN_WARN("batch submit: missing attachment {} — falling back to individual submit", NeedHash);
+					return FallbackToIndividualSubmit(Actions);
+				}
+			}
+
+			HttpClient::Response RetryResponse = m_Http.Post(SubmitUrl, Pkg);
+
+			if (RetryResponse.StatusCode == HttpResponseCode::OK)
+			{
+				return ParseBatchResponse(RetryResponse, Actions);
+			}
+
+			ZEN_WARN("batch submit retry failed with {} {} — falling back to individual submit",
+					 (int)RetryResponse.StatusCode,
+					 ToString(RetryResponse.StatusCode));
+			return FallbackToIndividualSubmit(Actions);
+		}
 	}
 
 	// Unexpected status or connection error — fall back to individual submission
 
-	if (Response)
-	{
-		ZEN_WARN("batch submit to {}{} returned {} {} — falling back to individual submit",
-				 m_Http.GetBaseUri(),
-				 SubmitUrl,
-				 (int)Response.StatusCode,
-				 ToString(Response.StatusCode));
-	}
-	else
-	{
-		ZEN_WARN("batch submit to {}{} failed — falling back to individual submit", m_Http.GetBaseUri(), SubmitUrl);
-	}
+	ZEN_WARN("batch submit to {}{} failed — falling back to individual submit", m_Http.GetBaseUri(), SubmitUrl);
 
 	return FallbackToIndividualSubmit(Actions);
 }
@@ -869,9 +947,10 @@ RemoteHttpRunner::SweepRunningActions()
 	{
 		for (auto& FieldIt : Completed["completed"sv])
 		{
-			CbObjectView	 EntryObj	 = FieldIt.AsObjectView();
-			const int32_t	 CompleteLsn = EntryObj["lsn"sv].AsInt32();
-			std::string_view StateName	 = EntryObj["state"sv].AsString();
+			CbObjectView	 EntryObj	   = FieldIt.AsObjectView();
+			const int32_t	 CompleteLsn   = EntryObj["lsn"sv].AsInt32();
+			std::string_view StateName	   = EntryObj["state"sv].AsString();
+			std::string_view FailureReason = EntryObj["reason"sv].AsString();
 
 			RunnerAction::State RemoteState = RunnerAction::FromString(StateName);
 
@@ -884,6 +963,7 @@ RemoteHttpRunner::SweepRunningActions()
 				{
 					HttpRunningAction CompletedAction = std::move(CompleteIt->second);
 					CompletedAction.RemoteState		  = RemoteState;
+					CompletedAction.FailureReason	  = std::string(FailureReason);
 
 					if (RemoteState == RunnerAction::State::Completed && ResponseJob)
 					{
@@ -927,16 +1007,44 @@ RemoteHttpRunner::SweepRunningActions()
 	{
 		const int ActionLsn = HttpAction.Action->ActionLsn;
 
-		ZEN_DEBUG("action {} LSN {} (remote LSN {}) -> {}",
-				  HttpAction.Action->ActionId,
-				  ActionLsn,
-				  HttpAction.RemoteActionLsn,
-				  RunnerAction::ToString(HttpAction.RemoteState));
-
 		if (HttpAction.RemoteState == RunnerAction::State::Completed)
 		{
+			ZEN_DEBUG("action {} LSN {} (remote LSN {}) completed on {}",
+					  HttpAction.Action->ActionId,
+					  ActionLsn,
+					  HttpAction.RemoteActionLsn,
+					  m_HostName);
 			HttpAction.Action->SetResult(std::move(HttpAction.ActionResults));
 		}
+		else if (HttpAction.RemoteState == RunnerAction::State::Failed || HttpAction.RemoteState == RunnerAction::State::Abandoned)
+		{
+			HttpAction.Action->FailureReason = HttpAction.FailureReason;
+			if (HttpAction.FailureReason.empty())
+			{
+				ZEN_WARN("action {} ({}) {} on remote {}",
+						 HttpAction.Action->ActionId,
+						 ActionLsn,
+						 RunnerAction::ToString(HttpAction.RemoteState),
+						 m_HostName);
+			}
+			else
+			{
+				ZEN_WARN("action {} ({}) {} on remote {}: {}",
+						 HttpAction.Action->ActionId,
+						 ActionLsn,
+						 RunnerAction::ToString(HttpAction.RemoteState),
+						 m_HostName,
+						 HttpAction.FailureReason);
+			}
+		}
+		else
+		{
+			ZEN_DEBUG("action {} LSN {} (remote LSN {}) -> {}",
+					  HttpAction.Action->ActionId,
+					  ActionLsn,
+					  HttpAction.RemoteActionLsn,
+					  RunnerAction::ToString(HttpAction.RemoteState));
+		}
 
 		HttpAction.Action->SetActionState(HttpAction.RemoteState);
 	}