zcbot/windows-node/Zcbot.WindowsNode/NodeConnectionLoop.cs

531 lines
22 KiB
C#

using System.Net;
using System.Net.WebSockets;
using System.Reflection;
using System.Runtime.InteropServices;
using System.Text;
using System.Text.Json;
using System.Collections.Concurrent;
namespace Zcbot.WindowsNode;
internal sealed class NodeConnectionLoop(NodeConfig config, Action<NodeStatus>? statusChanged = null)
{
private readonly SemaphoreSlim sendLock = new(1, 1);
private readonly JobInboxStore jobInbox = new(NodePaths.ForCurrentMachine().JobsDirectory);
private readonly JobInputDownloader inputDownloader = new(
config, new JobInboxStore(NodePaths.ForCurrentMachine().JobsDirectory));
private readonly NodeAdapterRegistry adapters = NodeAdapterRegistry.CreateDefault(
new JobInboxStore(NodePaths.ForCurrentMachine().JobsDirectory));
private readonly JobOutputUploader outputUploader = new(config);
private readonly ConcurrentDictionary<Guid, Task> jobPipelines = new();
private static readonly TimeSpan[] Backoff =
[
TimeSpan.FromSeconds(1),
TimeSpan.FromSeconds(2),
TimeSpan.FromSeconds(5),
TimeSpan.FromSeconds(10),
TimeSpan.FromSeconds(30),
TimeSpan.FromSeconds(60),
];
internal async Task RunAsync(CancellationToken cancellationToken)
{
var attempt = 0;
while (!cancellationToken.IsCancellationRequested)
{
try
{
Report(NodeState.Connecting, "正在连接 zcbot…");
await ConnectOnceAsync(cancellationToken);
attempt = 0;
}
catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested)
{
throw;
}
catch (WebSocketException exception)
{
Report(NodeState.Offline, "连接中断,等待重连");
Console.Error.WriteLine($"[WARN] WebSocket disconnected: {exception.Message}");
}
catch (IOException exception)
{
Report(NodeState.Offline, "网络不可用,等待重连");
Console.Error.WriteLine($"[WARN] Connection I/O failed: {exception.Message}");
}
catch (JsonException exception)
{
Report(NodeState.Offline, "服务端消息无效,等待重连");
Console.Error.WriteLine($"[WARN] Invalid server message: {exception.Message}");
}
catch (NodeEndpointException exception)
{
Report(NodeState.Offline, "WebSocket 握手被拒绝,请检查服务端或反向代理");
Console.Error.WriteLine($"[WARN] WebSocket handshake rejected: {exception.Message}");
}
var baseDelay = Backoff[Math.Min(attempt, Backoff.Length - 1)];
attempt++;
var jitter = TimeSpan.FromMilliseconds(Random.Shared.Next(0, 750));
var delay = baseDelay + jitter;
Console.WriteLine($"[INFO] Reconnecting in {delay.TotalSeconds:F1}s.");
await Task.Delay(delay, cancellationToken);
}
}
private async Task ConnectOnceAsync(CancellationToken cancellationToken)
{
using var socket = new ClientWebSocket();
socket.Options.SetRequestHeader("Authorization", $"Bearer {config.NodeToken}");
socket.Options.SetRequestHeader("X-Node-Id", config.NodeId.ToString());
socket.Options.KeepAliveInterval = TimeSpan.FromSeconds(config.HeartbeatSeconds);
socket.Options.CollectHttpResponseDetails = true;
var endpoint = NodeUri.WebSocketEndpoint(config.ServerUrl);
Console.WriteLine($"[INFO] Connecting to {endpoint.GetLeftPart(UriPartial.Path)}.");
try
{
await socket.ConnectAsync(endpoint, cancellationToken);
}
catch (WebSocketException exception) when (
socket.HttpStatusCode is HttpStatusCode.Unauthorized or HttpStatusCode.Forbidden)
{
throw new NodeEndpointException(
$"WebSocket handshake returned HTTP {(int?)socket.HttpStatusCode}. "
+ "Verify that the server includes the current Windows Node routes and that "
+ $"the reverse proxy forwards WebSocket Upgrade for {endpoint.AbsolutePath}. "
+ exception.Message);
}
Console.WriteLine("[OK] Node connected.");
Report(NodeState.Online, "已连接");
await SendAsync(socket, "hello", RuntimePayload(), cancellationToken);
await ReportRecoverableJobsAsync(socket, cancellationToken);
using var heartbeatStop = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken);
var heartbeat = HeartbeatLoopAsync(socket, heartbeatStop.Token);
try
{
await ReceiveLoopAsync(socket, cancellationToken);
}
finally
{
heartbeatStop.Cancel();
try
{
await heartbeat;
}
catch (OperationCanceledException) when (heartbeatStop.IsCancellationRequested)
{
}
}
}
private async Task HeartbeatLoopAsync(ClientWebSocket socket, CancellationToken cancellationToken)
{
using var timer = new PeriodicTimer(TimeSpan.FromSeconds(config.HeartbeatSeconds));
while (await timer.WaitForNextTickAsync(cancellationToken))
{
await SendAsync(socket, "heartbeat", RuntimePayload(), cancellationToken);
foreach (var job in jobInbox.ReadRecoverableJobs())
{
if (job.Terminal is JsonElement terminal
&& terminal.GetProperty("status").GetString() == "succeeded"
&& !job.UploadComplete)
{
StartJobPipeline(socket, job);
}
}
}
}
private async Task ReportRecoverableJobsAsync(
ClientWebSocket socket, CancellationToken cancellationToken)
{
foreach (var job in jobInbox.ReadRecoverableJobs())
{
if (job.Terminal is JsonElement terminal)
{
if (terminal.GetProperty("status").GetString() == "succeeded")
{
if (!job.UploadComplete)
{
StartJobPipeline(socket, job);
}
}
else
{
await SendAsync(socket, "job_terminal", terminal, cancellationToken);
}
continue;
}
await SendAsync(socket, "job_state", new
{
job_id = job.JobId,
lease_id = job.LeaseId,
request_digest = job.RequestDigest,
stage = "downloading_inputs",
progress = 0,
metrics = new { },
}, cancellationToken);
StartJobPipeline(socket, job);
}
}
private async Task ReceiveLoopAsync(
ClientWebSocket socket, CancellationToken cancellationToken)
{
var buffer = new byte[16 * 1024];
using var message = new MemoryStream();
while (socket.State == WebSocketState.Open)
{
var result = await socket.ReceiveAsync(buffer, cancellationToken);
if (result.MessageType == WebSocketMessageType.Close)
{
if ((int?)result.CloseStatus == 4003)
{
throw new NodeConfigurationException(
"节点身份已被服务端拒绝。请在管理后台确认该 Node ID 未被禁用或删除;"
+ "若记录不存在或身份已撤销,请清除本机身份并使用新注册码重新注册。"
+ $" 服务端信息:{result.CloseStatusDescription}");
}
if (result.CloseStatus == WebSocketCloseStatus.PolicyViolation)
{
throw new NodeConfigurationException(
$"节点上报被服务端拒绝:{result.CloseStatusDescription}");
}
return;
}
if (result.MessageType != WebSocketMessageType.Text)
{
throw new JsonException("Only text WebSocket messages are supported.");
}
message.Write(buffer, 0, result.Count);
if (!result.EndOfMessage)
{
if (message.Length > 1024 * 1024)
{
throw new JsonException("Server message exceeded 1 MiB.");
}
continue;
}
using var document = JsonDocument.Parse(message.ToArray());
if (document.RootElement.TryGetProperty("type", out var type))
{
Console.WriteLine($"[INFO] Server message: {type.GetString()}.");
if (type.GetString() == "job_offer"
&& document.RootElement.TryGetProperty("payload", out var payload))
{
var offerResult = jobInbox.Accept(payload, adapters);
await SendAsync(
socket,
offerResult.Accepted ? "job_accept" : "job_reject",
offerResult.Accepted
? new
{
job_id = offerResult.JobId,
lease_id = offerResult.LeaseId,
request_digest = offerResult.RequestDigest,
}
: new
{
job_id = payload.TryGetProperty("job_id", out var jobId)
? jobId.GetString() : "",
lease_id = payload.TryGetProperty("lease_id", out var leaseId)
? leaseId.GetString() : "",
reason = offerResult.Reason,
},
cancellationToken);
if (offerResult.Accepted)
{
await SendAsync(socket, "job_state", new
{
job_id = offerResult.JobId,
lease_id = offerResult.LeaseId,
request_digest = offerResult.RequestDigest,
stage = "downloading_inputs",
progress = 0,
metrics = new { },
}, cancellationToken);
var acceptedJob = jobInbox.ReadRecoverableJobs()
.Single(item => item.JobId == offerResult.JobId);
StartJobPipeline(socket, acceptedJob);
}
}
else if (type.GetString() == "job_cancel"
&& document.RootElement.TryGetProperty("payload", out var cancelPayload)
&& TryCancelJob(cancelPayload, out var cancelledJob))
{
adapters.Find(cancelledJob.Capability)?.Cancel(cancelledJob.JobId);
jobInbox.WriteTerminal(
cancelledJob, "cancelled", "USER_CANCELLED", "Cancelled by user.");
await SendAsync(socket, "job_terminal", new
{
job_id = cancelledJob.JobId,
lease_id = cancelledJob.LeaseId,
request_digest = cancelledJob.RequestDigest,
status = "cancelled",
error = new { code = "USER_CANCELLED", detail = "Cancelled by user." },
artifact_manifest = Array.Empty<object>(),
}, cancellationToken);
}
}
message.SetLength(0);
}
}
private bool TryCancelJob(JsonElement payload, out RecoverableJob job)
{
job = null!;
if (!payload.TryGetProperty("job_id", out var jobIdValue)
|| !Guid.TryParse(jobIdValue.GetString(), out var jobId)
|| !payload.TryGetProperty("lease_id", out var leaseIdValue)
|| !Guid.TryParse(leaseIdValue.GetString(), out var leaseId)
|| !payload.TryGetProperty("request_digest", out var digestValue))
{
return false;
}
job = jobInbox.ReadRecoverableJobs().SingleOrDefault(item =>
item.JobId == jobId
&& item.LeaseId == leaseId
&& item.RequestDigest == digestValue.GetString()
&& item.Terminal is null)!;
return job is not null;
}
private void StartJobPipeline(ClientWebSocket socket, RecoverableJob job)
{
var completion = new TaskCompletionSource(
TaskCreationOptions.RunContinuationsAsynchronously);
if (jobPipelines.TryAdd(job.JobId, completion.Task))
{
_ = RunJobPipelineAndReleaseAsync(socket, job, completion);
}
}
private async Task RunJobPipelineAndReleaseAsync(
ClientWebSocket socket, RecoverableJob job, TaskCompletionSource completion)
{
try
{
await RunJobPipelineAsync(socket, job);
}
finally
{
completion.TrySetResult();
jobPipelines.TryRemove(job.JobId, out _);
}
}
private async Task RunJobPipelineAsync(ClientWebSocket socket, RecoverableJob job)
{
var adapter = adapters.Find(job.Capability)
?? throw new InvalidDataException($"No local adapter for {job.Capability}.");
var recoveringOutputs = job.Terminal is JsonElement;
try
{
if (job.Terminal is null)
{
jobInbox.WriteState(job, "downloading_inputs", 0, "正在下载并校验输入文件");
await inputDownloader.DownloadAsync(job, CancellationToken.None);
var afterDownload = jobInbox.ReadRecoverableJobs()
.Single(item => item.JobId == job.JobId);
if (afterDownload.Terminal is JsonElement cancelledTerminal)
{
await TrySendAsync(socket, "job_terminal", cancelledTerminal);
return;
}
jobInbox.WriteState(job, "ready_to_run", 5, "输入文件已就绪,准备启动软件");
await TrySendAsync(socket, "job_state", new
{
job_id = job.JobId,
lease_id = job.LeaseId,
request_digest = job.RequestDigest,
stage = "ready_to_run",
progress = 5,
metrics = new
{
input_count = jobInbox.InputTransfers(job).Count,
input_bytes = jobInbox.InputTransfers(job)
.Sum(item => item.GetProperty("size_bytes").GetInt64()),
},
});
jobInbox.WriteState(job, "software_running", 10, adapter.RunningDetail);
await TrySendAsync(socket, "job_state", new
{
job_id = job.JobId,
lease_id = job.LeaseId,
request_digest = job.RequestDigest,
stage = "software_running",
progress = 10,
metrics = new { },
});
await adapter.RunAsync(job);
}
var refreshed = jobInbox.ReadRecoverableJobs()
.Single(item => item.JobId == job.JobId);
if (refreshed.Terminal is not JsonElement terminal)
{
throw new InvalidDataException("Software adapter did not create a terminal record.");
}
if (terminal.GetProperty("status").GetString() != "succeeded")
{
var terminalStatus = terminal.GetProperty("status").GetString() ?? "failed";
var terminalDetail = terminal.TryGetProperty("error", out var error)
&& error.TryGetProperty("detail", out var errorDetail)
? errorDetail.GetString() ?? "任务执行失败"
: "任务执行失败";
jobInbox.WriteState(
refreshed, terminalStatus, terminalStatus == "cancelled" ? 0 : 10, terminalDetail);
await TrySendAsync(socket, "job_terminal", terminal);
return;
}
jobInbox.WriteState(refreshed, "uploading_outputs", 90, "软件执行完成,正在上传结果");
await TrySendAsync(socket, "job_state", new
{
job_id = job.JobId,
lease_id = job.LeaseId,
request_digest = job.RequestDigest,
stage = "uploading_outputs",
progress = 90,
metrics = new { },
});
await outputUploader.UploadAsync(refreshed, recoveringOutputs);
jobInbox.WriteState(refreshed, "succeeded", 100, "结果已上传并由云端确认");
}
catch (Exception exception) when (
exception is HttpRequestException
or IOException
or JsonException
or UnauthorizedAccessException
or InvalidDataException)
{
var current = jobInbox.ReadRecoverableJobs()
.Single(item => item.JobId == job.JobId);
if (current.Terminal is JsonElement terminal
&& terminal.GetProperty("status").GetString() == "succeeded")
{
jobInbox.WriteState(
current,
"uploading_outputs",
90,
$"结果上传暂缓:{exception.Message}");
outputUploader.RecordDeferred(current, exception);
return;
}
jobInbox.WriteTerminal(
job,
"failed",
"INPUT_DOWNLOAD_FAILED",
exception.Message[..Math.Min(exception.Message.Length, 500)]);
jobInbox.WriteState(job, "failed", 0, exception.Message);
var failedTerminal = jobInbox.ReadRecoverableJobs()
.Single(item => item.JobId == job.JobId).Terminal;
if (failedTerminal is JsonElement payload)
{
await TrySendAsync(socket, "job_terminal", payload);
}
}
}
private async Task TrySendTerminalAsync(ClientWebSocket socket, RecoverableJob job)
{
var terminal = jobInbox.ReadRecoverableJobs()
.Single(item => item.JobId == job.JobId).Terminal;
if (terminal is JsonElement terminalPayload)
{
await TrySendAsync(socket, "job_terminal", terminalPayload);
}
}
private async Task TrySendAsync(ClientWebSocket socket, string type, object payload)
{
if (socket.State != WebSocketState.Open)
{
return;
}
try
{
await SendAsync(socket, type, payload, CancellationToken.None);
}
catch (Exception exception) when (
exception is WebSocketException or IOException or ObjectDisposedException)
{
Console.Error.WriteLine($"[WARN] Job report deferred: {exception.Message}");
}
}
private void Report(NodeState state, string message) =>
statusChanged?.Invoke(NodeStatus.Create(state, message));
private async Task SendAsync(
ClientWebSocket socket, string type, object payload, CancellationToken cancellationToken)
{
var envelope = JsonSerializer.SerializeToUtf8Bytes(new
{
protocol_version = 1,
message_id = Guid.NewGuid(),
type,
sent_at = DateTimeOffset.UtcNow,
payload,
});
await sendLock.WaitAsync(cancellationToken);
try
{
await socket.SendAsync(
envelope, WebSocketMessageType.Text, endOfMessage: true, cancellationToken);
}
finally
{
sendLock.Release();
}
}
private object RuntimePayload()
{
var root = Path.GetPathRoot(Environment.SystemDirectory) ?? "C:\\";
var origin = OriginRuntimeProbe.Detect();
var capabilityRuntime = adapters.All
.Where(item => config.Capabilities.Contains(item.Capability, StringComparer.Ordinal))
.ToDictionary(
item => item.Capability,
item =>
{
var runtime = item.DetectRuntime();
var slots = runtime.Health == "ready"
&& !jobInbox.HasPendingJobs
&& !item.HasActiveJobs ? 1 : 0;
return new
{
adapter_version = item.AdapterVersion,
features = item.Features,
available_slots = slots,
health = runtime.Health,
detail = runtime.Detail,
};
});
var originAdapter = adapters.Find(OriginPlotNodeAdapter.CapabilityName);
return new
{
install_id = config.InstallId,
node_name = config.NodeName,
node_version = Assembly.GetExecutingAssembly().GetName().Version?.ToString(3) ?? "0.1.0",
os_version = RuntimeInformation.OSDescription,
capabilities = config.Capabilities,
available_slots = origin.Health == "ready"
&& !jobInbox.HasPendingJobs
&& !(originAdapter?.HasActiveJobs ?? false) ? 1 : 0,
capability_runtime = capabilityRuntime,
disk_free_bytes = new DriveInfo(root).AvailableFreeSpace,
desktop_session = Environment.UserInteractive,
origin = new
{
software = origin.Software,
software_version = origin.SoftwareVersion,
adapter_version = origin.AdapterVersion,
health = origin.Health,
detail = origin.Detail,
},
};
}
}