484 lines
19 KiB
C#
484 lines
19 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 OriginWorkerRunner workerRunner = new(
|
|
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 = "waiting_input",
|
|
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);
|
|
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 = "waiting_input",
|
|
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))
|
|
{
|
|
workerRunner.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)
|
|
{
|
|
try
|
|
{
|
|
if (job.Terminal is null)
|
|
{
|
|
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;
|
|
}
|
|
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_bytes = job.InputTransfer?.GetProperty("size_bytes").GetInt64() },
|
|
});
|
|
await TrySendAsync(socket, "job_state", new
|
|
{
|
|
job_id = job.JobId,
|
|
lease_id = job.LeaseId,
|
|
request_digest = job.RequestDigest,
|
|
stage = "origin_running",
|
|
progress = 10,
|
|
metrics = new { },
|
|
});
|
|
await workerRunner.RunAsync(job);
|
|
}
|
|
var refreshed = jobInbox.ReadRecoverableJobs()
|
|
.Single(item => item.JobId == job.JobId);
|
|
if (refreshed.Terminal is not JsonElement terminal)
|
|
{
|
|
throw new InvalidDataException("Origin worker did not create a terminal record.");
|
|
}
|
|
if (terminal.GetProperty("status").GetString() != "succeeded")
|
|
{
|
|
await TrySendAsync(socket, "job_terminal", terminal);
|
|
return;
|
|
}
|
|
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);
|
|
}
|
|
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")
|
|
{
|
|
Console.Error.WriteLine($"[WARN] Output upload deferred: {exception.Message}");
|
|
return;
|
|
}
|
|
jobInbox.WriteTerminal(
|
|
job,
|
|
"failed",
|
|
"INPUT_DOWNLOAD_FAILED",
|
|
exception.Message[..Math.Min(exception.Message.Length, 500)]);
|
|
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();
|
|
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.HasPendingOriginJobs
|
|
&& !workerRunner.HasActiveJobs ? 1 : 0,
|
|
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,
|
|
},
|
|
};
|
|
}
|
|
}
|