feat: fence checker facts by lease token
Some checks are pending
ci / proto (push) Waiting to run
ci / test (ubuntu-latest) (push) Waiting to run
ci / test (windows-latest) (push) Waiting to run
ci / race (push) Waiting to run
ci / integration (push) Waiting to run

This commit is contained in:
youfak 2026-07-31 21:16:23 +08:00
parent 2bdc1ebda3
commit 9934c659ab
9 changed files with 81 additions and 30 deletions

View File

@ -235,6 +235,7 @@ message CheckTask {
string password = 14; string password = 14;
string credential_version = 15; string credential_version = 15;
uint32 max_attempts = 16; uint32 max_attempts = 16;
string lease_token = 17;
} }
enum CheckLevel { enum CheckLevel {
@ -260,6 +261,7 @@ message HealthObservation {
google.protobuf.Duration latency = 8; google.protobuf.Duration latency = 8;
string observed_egress_ip = 9; string observed_egress_ip = 9;
google.protobuf.Timestamp observed_at = 10; google.protobuf.Timestamp observed_at = 10;
string lease_token = 11;
} }
message ReportObservationsResponse { message ReportObservationsResponse {

View File

@ -171,6 +171,9 @@ Routing 决定 AVAILABLE、SUSPECT 或 UNHEALTHY并更新 Redis 活动池,
检查层级、目标 Profile 与领取者一致Reducer 成功后才确认任务;相同领取者对已确认 检查层级、目标 Profile 与领取者一致Reducer 成功后才确认任务;相同领取者对已确认
任务的同一事实可重放,由活动池摘要幂等处理。当前生产配置尚未装配 Redis 共享 broker 任务的同一事实可重放,由活动池摘要幂等处理。当前生产配置尚未装配 Redis 共享 broker
因此无共享 broker 的服务会以 `Unavailable` 拒绝任务流,而不下发无租约任务。 因此无共享 broker 的服务会以 `Unavailable` 拒绝任务流,而不下发无租约任务。
每次 Claim 还会签发新的不可预测 `lease_token`Observation 必须回传该值。任务被重新
领取后,旧 token 即使拥有相同 `task_id``checker_id` 也会被拒绝,避免过期实例的
迟到事实覆盖新租约结果。
`proxy-checker` 使用固定大小 worker-pool 执行每个 pull 批次,任务数不超过该请求的 `proxy-checker` 使用固定大小 worker-pool 执行每个 pull 批次,任务数不超过该请求的
`max_in_flight`;每次尝试都受 `deadline``timeout` 的较小值约束,失败可在同一 `max_in_flight`;每次尝试都受 `deadline``timeout` 的较小值约束,失败可在同一

View File

@ -281,7 +281,7 @@ Lua 同一原子边界归并,覆盖首次 CHECKING 失败进入 UNHEALTHY、AV
Checker Observation 上报 RPC 已复用既有控制面监听接入 ControllerSPIFFE `checker` 身份、 Checker Observation 上报 RPC 已复用既有控制面监听接入 ControllerSPIFFE `checker` 身份、
每批上限、配置阈值解析、代理归属查询与 Reducer 均已闭环。`StreamCheckTasks` 现已实现为 每批上限、配置阈值解析、代理归属查询与 Reducer 均已闭环。`StreamCheckTasks` 现已实现为
有界 pull并通过通用任务 broker 契约完成能力协商、同 Checker 并发窗口、租约到期回收、 有界 pull并通过通用任务 broker 契约完成能力协商、同 Checker 并发窗口、租约到期回收、
领取者栅栏和完成后重放任务凭据仅由认证流在执行期下发。Redis 共享 due-index/租约持久化、 领取者加不可预测 lease token 的栅栏和完成后重放任务凭据仅由认证流在执行期下发。Redis 共享 due-index/租约持久化、
生产 broker 已有契约但 Redis 实现尚未完成;`proxy-checker` 独立进程、固定大小 worker-pool、 生产 broker 已有契约但 Redis 实现尚未完成;`proxy-checker` 独立进程、固定大小 worker-pool、
任务期重试/微批上报和 HTTP/HTTPS BASIC/TARGET 探测器已完成并有测试。SOCKS5、EGRESS 任务期重试/微批上报和 HTTP/HTTPS BASIC/TARGET 探测器已完成并有测试。SOCKS5、EGRESS
专用出口探测、Redis due-index/跨副本租约和部署运行态仍未实现,因此本任务保持未完成。 专用出口探测、Redis due-index/跨副本租约和部署运行态仍未实现,因此本任务保持未完成。

View File

@ -1801,6 +1801,7 @@ type CheckTask struct {
Password string `protobuf:"bytes,14,opt,name=password,proto3" json:"password,omitempty"` Password string `protobuf:"bytes,14,opt,name=password,proto3" json:"password,omitempty"`
CredentialVersion string `protobuf:"bytes,15,opt,name=credential_version,json=credentialVersion,proto3" json:"credential_version,omitempty"` CredentialVersion string `protobuf:"bytes,15,opt,name=credential_version,json=credentialVersion,proto3" json:"credential_version,omitempty"`
MaxAttempts uint32 `protobuf:"varint,16,opt,name=max_attempts,json=maxAttempts,proto3" json:"max_attempts,omitempty"` MaxAttempts uint32 `protobuf:"varint,16,opt,name=max_attempts,json=maxAttempts,proto3" json:"max_attempts,omitempty"`
LeaseToken string `protobuf:"bytes,17,opt,name=lease_token,json=leaseToken,proto3" json:"lease_token,omitempty"`
unknownFields protoimpl.UnknownFields unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache sizeCache protoimpl.SizeCache
} }
@ -1947,6 +1948,13 @@ func (x *CheckTask) GetMaxAttempts() uint32 {
return 0 return 0
} }
func (x *CheckTask) GetLeaseToken() string {
if x != nil {
return x.LeaseToken
}
return ""
}
type ObservationBatch struct { type ObservationBatch struct {
state protoimpl.MessageState `protogen:"open.v1"` state protoimpl.MessageState `protogen:"open.v1"`
CheckerId string `protobuf:"bytes,1,opt,name=checker_id,json=checkerId,proto3" json:"checker_id,omitempty"` CheckerId string `protobuf:"bytes,1,opt,name=checker_id,json=checkerId,proto3" json:"checker_id,omitempty"`
@ -2011,6 +2019,7 @@ type HealthObservation struct {
Latency *durationpb.Duration `protobuf:"bytes,8,opt,name=latency,proto3" json:"latency,omitempty"` Latency *durationpb.Duration `protobuf:"bytes,8,opt,name=latency,proto3" json:"latency,omitempty"`
ObservedEgressIp string `protobuf:"bytes,9,opt,name=observed_egress_ip,json=observedEgressIp,proto3" json:"observed_egress_ip,omitempty"` ObservedEgressIp string `protobuf:"bytes,9,opt,name=observed_egress_ip,json=observedEgressIp,proto3" json:"observed_egress_ip,omitempty"`
ObservedAt *timestamppb.Timestamp `protobuf:"bytes,10,opt,name=observed_at,json=observedAt,proto3" json:"observed_at,omitempty"` ObservedAt *timestamppb.Timestamp `protobuf:"bytes,10,opt,name=observed_at,json=observedAt,proto3" json:"observed_at,omitempty"`
LeaseToken string `protobuf:"bytes,11,opt,name=lease_token,json=leaseToken,proto3" json:"lease_token,omitempty"`
unknownFields protoimpl.UnknownFields unknownFields protoimpl.UnknownFields
sizeCache protoimpl.SizeCache sizeCache protoimpl.SizeCache
} }
@ -2115,6 +2124,13 @@ func (x *HealthObservation) GetObservedAt() *timestamppb.Timestamp {
return nil return nil
} }
func (x *HealthObservation) GetLeaseToken() string {
if x != nil {
return x.LeaseToken
}
return ""
}
type ReportObservationsResponse struct { type ReportObservationsResponse struct {
state protoimpl.MessageState `protogen:"open.v1"` state protoimpl.MessageState `protogen:"open.v1"`
Accepted uint32 `protobuf:"varint,1,opt,name=accepted,proto3" json:"accepted,omitempty"` Accepted uint32 `protobuf:"varint,1,opt,name=accepted,proto3" json:"accepted,omitempty"`
@ -2324,7 +2340,7 @@ const file_controlplane_v1_controlplane_proto_rawDesc = "" +
"\vinstance_id\x18\x02 \x01(\tR\n" + "\vinstance_id\x18\x02 \x01(\tR\n" +
"instanceId\x12\"\n" + "instanceId\x12\"\n" +
"\rmax_in_flight\x18\x03 \x01(\rR\vmaxInFlight\x12Q\n" + "\rmax_in_flight\x18\x03 \x01(\rR\vmaxInFlight\x12Q\n" +
"\x10supported_levels\x18\x04 \x03(\x0e2&.proxy_pool.controlplane.v1.CheckLevelR\x0fsupportedLevels\"\xde\x04\n" + "\x10supported_levels\x18\x04 \x03(\x0e2&.proxy_pool.controlplane.v1.CheckLevelR\x0fsupportedLevels\"\xff\x04\n" +
"\tCheckTask\x12\x17\n" + "\tCheckTask\x12\x17\n" +
"\atask_id\x18\x01 \x01(\tR\x06taskId\x12\x19\n" + "\atask_id\x18\x01 \x01(\tR\x06taskId\x12\x19\n" +
"\bproxy_id\x18\x02 \x01(\tR\aproxyId\x12E\n" + "\bproxy_id\x18\x02 \x01(\tR\aproxyId\x12E\n" +
@ -2344,11 +2360,13 @@ const file_controlplane_v1_controlplane_proto_rawDesc = "" +
"\busername\x18\r \x01(\tR\busername\x12\x1a\n" + "\busername\x18\r \x01(\tR\busername\x12\x1a\n" +
"\bpassword\x18\x0e \x01(\tR\bpassword\x12-\n" + "\bpassword\x18\x0e \x01(\tR\bpassword\x12-\n" +
"\x12credential_version\x18\x0f \x01(\tR\x11credentialVersion\x12!\n" + "\x12credential_version\x18\x0f \x01(\tR\x11credentialVersion\x12!\n" +
"\fmax_attempts\x18\x10 \x01(\rR\vmaxAttempts\"\x84\x01\n" + "\fmax_attempts\x18\x10 \x01(\rR\vmaxAttempts\x12\x1f\n" +
"\vlease_token\x18\x11 \x01(\tR\n" +
"leaseToken\"\x84\x01\n" +
"\x10ObservationBatch\x12\x1d\n" + "\x10ObservationBatch\x12\x1d\n" +
"\n" + "\n" +
"checker_id\x18\x01 \x01(\tR\tcheckerId\x12Q\n" + "checker_id\x18\x01 \x01(\tR\tcheckerId\x12Q\n" +
"\fobservations\x18\x02 \x03(\v2-.proxy_pool.controlplane.v1.HealthObservationR\fobservations\"\xa6\x03\n" + "\fobservations\x18\x02 \x03(\v2-.proxy_pool.controlplane.v1.HealthObservationR\fobservations\"\xc7\x03\n" +
"\x11HealthObservation\x12\x17\n" + "\x11HealthObservation\x12\x17\n" +
"\atask_id\x18\x01 \x01(\tR\x06taskId\x12\x19\n" + "\atask_id\x18\x01 \x01(\tR\x06taskId\x12\x19\n" +
"\bproxy_id\x18\x02 \x01(\tR\aproxyId\x12<\n" + "\bproxy_id\x18\x02 \x01(\tR\aproxyId\x12<\n" +
@ -2362,7 +2380,9 @@ const file_controlplane_v1_controlplane_proto_rawDesc = "" +
"\x12observed_egress_ip\x18\t \x01(\tR\x10observedEgressIp\x12;\n" + "\x12observed_egress_ip\x18\t \x01(\tR\x10observedEgressIp\x12;\n" +
"\vobserved_at\x18\n" + "\vobserved_at\x18\n" +
" \x01(\v2\x1a.google.protobuf.TimestampR\n" + " \x01(\v2\x1a.google.protobuf.TimestampR\n" +
"observedAt\"T\n" + "observedAt\x12\x1f\n" +
"\vlease_token\x18\v \x01(\tR\n" +
"leaseToken\"T\n" +
"\x1aReportObservationsResponse\x12\x1a\n" + "\x1aReportObservationsResponse\x12\x1a\n" +
"\baccepted\x18\x01 \x01(\rR\baccepted\x12\x1a\n" + "\baccepted\x18\x01 \x01(\rR\baccepted\x12\x1a\n" +
"\brejected\x18\x02 \x01(\rR\brejected*\xc5\x01\n" + "\brejected\x18\x02 \x01(\rR\brejected*\xc5\x01\n" +

View File

@ -212,7 +212,7 @@ func (runner *Runner) executeTask(ctx context.Context, task *controlplanev1.Chec
result.FailureClass = probe.FailureProxyRequest result.FailureClass = probe.FailureProxyRequest
} }
observation := &controlplanev1.HealthObservation{ observation := &controlplanev1.HealthObservation{
TaskId: task.GetTaskId(), ProxyId: task.GetProxyId(), Level: task.GetLevel(), Success: result.Success, TaskId: task.GetTaskId(), LeaseToken: task.GetLeaseToken(), ProxyId: task.GetProxyId(), Level: task.GetLevel(), Success: result.Success,
FailureClass: result.FailureClass, Latency: durationpb.New(result.Latency), FailureClass: result.FailureClass, Latency: durationpb.New(result.Latency),
ObservedAt: timestamppb.New(runner.options.Now().UTC()), ObservedAt: timestamppb.New(runner.options.Now().UTC()),
} }

View File

@ -146,7 +146,7 @@ func (handler *GRPCHandler) ReportObservations(
if now.IsZero() { if now.IsZero() {
return nil, healthGRPCError(ErrInvalidGRPCHandler) return nil, healthGRPCError(ErrInvalidGRPCHandler)
} }
err = handler.options.TaskBroker.AuthorizeObservation(ctx, request.GetCheckerId(), observation, now) err = handler.options.TaskBroker.AuthorizeObservation(ctx, request.GetCheckerId(), item.GetLeaseToken(), observation, now)
} }
if err == nil { if err == nil {
_, err = handler.reducer.Apply(ctx, observation) _, err = handler.reducer.Apply(ctx, observation)
@ -156,7 +156,7 @@ func (handler *GRPCHandler) ReportObservations(
if now.IsZero() { if now.IsZero() {
return nil, healthGRPCError(ErrInvalidGRPCHandler) return nil, healthGRPCError(ErrInvalidGRPCHandler)
} }
err = handler.options.TaskBroker.CompleteObservation(ctx, request.GetCheckerId(), observation, now) err = handler.options.TaskBroker.CompleteObservation(ctx, request.GetCheckerId(), item.GetLeaseToken(), observation, now)
} }
if err == nil { if err == nil {
response.Accepted++ response.Accepted++
@ -226,7 +226,7 @@ func healthGRPCError(err error) error {
} }
func wireCheckTask(task LeasedTask, now time.Time) (*controlplanev1.CheckTask, error) { func wireCheckTask(task LeasedTask, now time.Time) (*controlplanev1.CheckTask, error) {
if !validTaskIdentifier(task.TaskID) || !validTaskIdentifier(task.ProxyID) || task.Host == "" || if !validTaskIdentifier(task.TaskID) || !validTaskIdentifier(task.LeaseToken) || !validTaskIdentifier(task.ProxyID) || task.Host == "" ||
strings.TrimSpace(task.Host) != task.Host || task.Port == 0 || task.Attempts <= 0 || task.Attempts > math.MaxUint32 || strings.TrimSpace(task.Host) != task.Host || task.Port == 0 || task.Attempts <= 0 || task.Attempts > math.MaxUint32 ||
!task.Deadline.After(now) || (task.SecretRef == "") != (task.CredentialVersion == "") { !task.Deadline.After(now) || (task.SecretRef == "") != (task.CredentialVersion == "") {
return nil, ErrInvalidLeasedTask return nil, ErrInvalidLeasedTask
@ -249,7 +249,7 @@ func wireCheckTask(task LeasedTask, now time.Time) (*controlplanev1.CheckTask, e
return nil, ErrInvalidLeasedTask return nil, ErrInvalidLeasedTask
} }
return &controlplanev1.CheckTask{ return &controlplanev1.CheckTask{
TaskId: task.TaskID, ProxyId: task.ProxyID, Protocol: protocol, Host: task.Host, Port: uint32(task.Port), TaskId: task.TaskID, LeaseToken: task.LeaseToken, ProxyId: task.ProxyID, Protocol: protocol, Host: task.Host, Port: uint32(task.Port),
SecretRef: task.SecretRef, CredentialVersion: task.CredentialVersion, Username: task.Username, Password: task.Password, SecretRef: task.SecretRef, CredentialVersion: task.CredentialVersion, Username: task.Username, Password: task.Password,
Level: level, RoutingName: task.RoutingName, TargetUrl: task.TargetURL, Level: level, RoutingName: task.RoutingName, TargetUrl: task.TargetURL,
Timeout: durationpb.New(task.Deadline.Sub(now)), Attempt: 1, MaxAttempts: uint32(task.Attempts), Timeout: durationpb.New(task.Deadline.Sub(now)), Attempt: 1, MaxAttempts: uint32(task.Attempts),

View File

@ -151,17 +151,21 @@ func TestGRPCHandlerStreamsLeasedTasksAndFencesReportedFacts(t *testing.T) {
if _, err := stream.Recv(); !errors.Is(err, io.EOF) { if _, err := stream.Recv(); !errors.Is(err, io.EOF) {
t.Fatalf("second Recv() error = %v, want EOF", err) t.Fatalf("second Recv() error = %v, want EOF", err)
} }
acceptedObservation := grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now.Add(time.Second))
acceptedObservation.LeaseToken = task.GetLeaseToken()
response, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{ response, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
CheckerId: "checker-a", Observations: []*controlplanev1.HealthObservation{ CheckerId: "checker-a", Observations: []*controlplanev1.HealthObservation{
grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now.Add(time.Second)), acceptedObservation,
}, },
}) })
if err != nil || response.GetAccepted() != 1 || len(global.commands) != 1 { if err != nil || response.GetAccepted() != 1 || len(global.commands) != 1 {
t.Fatalf("ReportObservations() = (%+v, %v), global=%d", response, err, len(global.commands)) t.Fatalf("ReportObservations() = (%+v, %v), global=%d", response, err, len(global.commands))
} }
rejectedObservation := grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now.Add(time.Second))
rejectedObservation.LeaseToken = task.GetLeaseToken()
rejected, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{ rejected, err := client.ReportObservations(context.Background(), &controlplanev1.ObservationBatch{
CheckerId: "checker-b", Observations: []*controlplanev1.HealthObservation{ CheckerId: "checker-b", Observations: []*controlplanev1.HealthObservation{
grpcHealthObservation(task.GetTaskId(), controlplanev1.CheckLevel_CHECK_LEVEL_BASIC, now.Add(time.Second)), rejectedObservation,
}, },
}) })
if err != nil || rejected.GetRejected() != 1 || len(global.commands) != 1 { if err != nil || rejected.GetRejected() != 1 || len(global.commands) != 1 {

View File

@ -2,6 +2,7 @@ package health
import ( import (
"context" "context"
"crypto/rand"
"crypto/sha256" "crypto/sha256"
"encoding/hex" "encoding/hex"
"errors" "errors"
@ -73,6 +74,7 @@ type TaskClaim struct {
// for the duration of a task lease and must not be logged. // for the duration of a task lease and must not be logged.
type LeasedTask struct { type LeasedTask struct {
TaskID string TaskID string
LeaseToken string
ProxyID string ProxyID string
Protocol proxyDomain.Scheme Protocol proxyDomain.Scheme
Host string Host string
@ -98,8 +100,8 @@ func (LeasedTask) Format(state fmt.State, _ rune) {
type TaskBroker interface { type TaskBroker interface {
TaskSink TaskSink
Claim(context.Context, TaskClaim) ([]LeasedTask, error) Claim(context.Context, TaskClaim) ([]LeasedTask, error)
AuthorizeObservation(context.Context, string, healthDomain.Observation, time.Time) error AuthorizeObservation(context.Context, string, string, healthDomain.Observation, time.Time) error
CompleteObservation(context.Context, string, healthDomain.Observation, time.Time) error CompleteObservation(context.Context, string, string, healthDomain.Observation, time.Time) error
} }
type MemoryTaskBrokerOptions struct { type MemoryTaskBrokerOptions struct {
@ -127,11 +129,13 @@ type brokerTask struct {
type taskLease struct { type taskLease struct {
checkerID string checkerID string
instanceID string instanceID string
token string
expiresAt time.Time expiresAt time.Time
} }
type taskCompletion struct { type taskCompletion struct {
checkerID string checkerID string
token string
expiresAt time.Time expiresAt time.Time
} }
@ -218,6 +222,7 @@ func (broker *MemoryTaskBroker) Claim(ctx context.Context, claim TaskClaim) ([]L
return []LeasedTask{}, nil return []LeasedTask{}, nil
} }
materials := make([]TaskMaterial, len(taskIDs)) materials := make([]TaskMaterial, len(taskIDs))
tokens := make([]string, len(taskIDs))
for index, taskID := range taskIDs { for index, taskID := range taskIDs {
material, resolveErr := broker.options.Material.ResolveCheckTask(ctx, broker.tasks[taskID].plan.Candidate) material, resolveErr := broker.options.Material.ResolveCheckTask(ctx, broker.tasks[taskID].plan.Candidate)
if resolveErr != nil { if resolveErr != nil {
@ -227,6 +232,11 @@ func (broker *MemoryTaskBroker) Claim(ctx context.Context, claim TaskClaim) ([]L
return nil, ErrInvalidLeasedTask return nil, ErrInvalidLeasedTask
} }
materials[index] = material materials[index] = material
leaseToken, tokenErr := newLeaseToken()
if tokenErr != nil {
return nil, tokenErr
}
tokens[index] = leaseToken
} }
result := make([]LeasedTask, len(taskIDs)) result := make([]LeasedTask, len(taskIDs))
for index, taskID := range taskIDs { for index, taskID := range taskIDs {
@ -235,8 +245,9 @@ func (broker *MemoryTaskBroker) Claim(ctx context.Context, claim TaskClaim) ([]L
if plan.Deadline.Before(leaseUntil) { if plan.Deadline.Before(leaseUntil) {
leaseUntil = plan.Deadline leaseUntil = plan.Deadline
} }
broker.leases[taskID] = taskLease{checkerID: claim.CheckerID, instanceID: claim.InstanceID, expiresAt: leaseUntil} leaseToken := tokens[index]
result[index] = leasedTask(taskID, plan, materials[index]) broker.leases[taskID] = taskLease{checkerID: claim.CheckerID, instanceID: claim.InstanceID, token: leaseToken, expiresAt: leaseUntil}
result[index] = leasedTask(taskID, leaseToken, plan, materials[index])
} }
return result, nil return result, nil
} }
@ -244,10 +255,11 @@ func (broker *MemoryTaskBroker) Claim(ctx context.Context, claim TaskClaim) ([]L
func (broker *MemoryTaskBroker) AuthorizeObservation( func (broker *MemoryTaskBroker) AuthorizeObservation(
ctx context.Context, ctx context.Context,
checkerID string, checkerID string,
leaseToken string,
observation healthDomain.Observation, observation healthDomain.Observation,
now time.Time, now time.Time,
) error { ) error {
if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || now.IsZero() { if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || !validTaskIdentifier(leaseToken) || now.IsZero() {
return ErrInvalidTaskBroker return ErrInvalidTaskBroker
} }
if err := ctx.Err(); err != nil { if err := ctx.Err(); err != nil {
@ -259,16 +271,17 @@ func (broker *MemoryTaskBroker) AuthorizeObservation(
} }
broker.mu.Lock() broker.mu.Lock()
defer broker.mu.Unlock() defer broker.mu.Unlock()
return broker.authorizeObservationLocked(checkerID, normalized, now.UTC()) return broker.authorizeObservationLocked(checkerID, leaseToken, normalized, now.UTC())
} }
func (broker *MemoryTaskBroker) CompleteObservation( func (broker *MemoryTaskBroker) CompleteObservation(
ctx context.Context, ctx context.Context,
checkerID string, checkerID string,
leaseToken string,
observation healthDomain.Observation, observation healthDomain.Observation,
now time.Time, now time.Time,
) error { ) error {
if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || now.IsZero() { if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || !validTaskIdentifier(leaseToken) || now.IsZero() {
return ErrInvalidTaskBroker return ErrInvalidTaskBroker
} }
if err := ctx.Err(); err != nil { if err := ctx.Err(); err != nil {
@ -281,7 +294,7 @@ func (broker *MemoryTaskBroker) CompleteObservation(
now = now.UTC() now = now.UTC()
broker.mu.Lock() broker.mu.Lock()
defer broker.mu.Unlock() defer broker.mu.Unlock()
if err := broker.authorizeObservationLocked(checkerID, normalized, now); err != nil { if err := broker.authorizeObservationLocked(checkerID, leaseToken, normalized, now); err != nil {
return err return err
} }
if _, complete := broker.finished[normalized.TaskID]; complete { if _, complete := broker.finished[normalized.TaskID]; complete {
@ -289,12 +302,13 @@ func (broker *MemoryTaskBroker) CompleteObservation(
} }
task := broker.tasks[normalized.TaskID] task := broker.tasks[normalized.TaskID]
delete(broker.leases, normalized.TaskID) delete(broker.leases, normalized.TaskID)
broker.finished[normalized.TaskID] = taskCompletion{checkerID: checkerID, expiresAt: task.plan.Deadline} broker.finished[normalized.TaskID] = taskCompletion{checkerID: checkerID, token: leaseToken, expiresAt: task.plan.Deadline}
return nil return nil
} }
func (broker *MemoryTaskBroker) authorizeObservationLocked( func (broker *MemoryTaskBroker) authorizeObservationLocked(
checkerID string, checkerID string,
leaseToken string,
observation healthDomain.Observation, observation healthDomain.Observation,
now time.Time, now time.Time,
) error { ) error {
@ -307,7 +321,7 @@ func (broker *MemoryTaskBroker) authorizeObservationLocked(
return ErrTaskObservation return ErrTaskObservation
} }
if finished, complete := broker.finished[observation.TaskID]; complete { if finished, complete := broker.finished[observation.TaskID]; complete {
if finished.checkerID != checkerID { if finished.checkerID != checkerID || finished.token != leaseToken {
return ErrTaskLeaseNotOwned return ErrTaskLeaseNotOwned
} }
return nil return nil
@ -316,7 +330,7 @@ func (broker *MemoryTaskBroker) authorizeObservationLocked(
if !leased || !lease.expiresAt.After(now) { if !leased || !lease.expiresAt.After(now) {
return ErrTaskLeaseExpired return ErrTaskLeaseExpired
} }
if lease.checkerID != checkerID { if lease.checkerID != checkerID || lease.token != leaseToken {
return ErrTaskLeaseNotOwned return ErrTaskLeaseNotOwned
} }
return nil return nil
@ -445,15 +459,23 @@ func deterministicTaskID(task PlannedTask) string {
return "check_" + hex.EncodeToString(digest[:]) return "check_" + hex.EncodeToString(digest[:])
} }
func leasedTask(taskID string, plan PlannedTask, material TaskMaterial) LeasedTask { func leasedTask(taskID, leaseToken string, plan PlannedTask, material TaskMaterial) LeasedTask {
return LeasedTask{ return LeasedTask{
TaskID: taskID, ProxyID: plan.Candidate.ProxyID, Protocol: material.Protocol, Host: material.Host, Port: material.Port, TaskID: taskID, LeaseToken: leaseToken, ProxyID: plan.Candidate.ProxyID, Protocol: material.Protocol, Host: material.Host, Port: material.Port,
SecretRef: material.SecretRef, CredentialVersion: material.CredentialVersion, Username: material.Username, Password: material.Password, SecretRef: material.SecretRef, CredentialVersion: material.CredentialVersion, Username: material.Username, Password: material.Password,
Level: plan.Candidate.Level, RoutingName: plan.Candidate.RoutingName, TargetURL: plan.Candidate.TargetURL, Level: plan.Candidate.Level, RoutingName: plan.Candidate.RoutingName, TargetURL: plan.Candidate.TargetURL,
Deadline: plan.Deadline.UTC(), Attempts: plan.Attempts, Deadline: plan.Deadline.UTC(), Attempts: plan.Attempts,
} }
} }
func newLeaseToken() (string, error) {
var entropy [24]byte
if _, err := rand.Read(entropy[:]); err != nil {
return "", err
}
return "lease_" + hex.EncodeToString(entropy[:]), nil
}
func taskMatchesObservation(task PlannedTask, observation healthDomain.Observation) bool { func taskMatchesObservation(task PlannedTask, observation healthDomain.Observation) bool {
return task.Candidate.ProxyID == observation.ProxyID && task.Candidate.Level == observation.Level && return task.Candidate.ProxyID == observation.ProxyID && task.Candidate.Level == observation.Level &&
task.Candidate.RoutingName == observation.RoutingName && task.Candidate.TargetURL == observation.TargetURL task.Candidate.RoutingName == observation.RoutingName && task.Candidate.TargetURL == observation.TargetURL

View File

@ -42,16 +42,16 @@ func TestMemoryTaskBrokerClaimsSupportedTasksAndFencesObservations(t *testing.T)
TaskID: claimed[0].TaskID, ProxyID: "basic", Level: healthDomain.LevelBasic, TaskID: claimed[0].TaskID, ProxyID: "basic", Level: healthDomain.LevelBasic,
Success: true, Latency: time.Millisecond, ObservedAt: now.Add(time.Second), Success: true, Latency: time.Millisecond, ObservedAt: now.Add(time.Second),
} }
if err := broker.AuthorizeObservation(context.Background(), "checker-b", observation, now.Add(time.Second)); !errors.Is(err, ErrTaskLeaseNotOwned) { if err := broker.AuthorizeObservation(context.Background(), "checker-b", claimed[0].LeaseToken, observation, now.Add(time.Second)); !errors.Is(err, ErrTaskLeaseNotOwned) {
t.Fatalf("AuthorizeObservation(other checker) error = %v, want ErrTaskLeaseNotOwned", err) t.Fatalf("AuthorizeObservation(other checker) error = %v, want ErrTaskLeaseNotOwned", err)
} }
if err := broker.AuthorizeObservation(context.Background(), "checker-a", observation, now.Add(time.Second)); err != nil { if err := broker.AuthorizeObservation(context.Background(), "checker-a", claimed[0].LeaseToken, observation, now.Add(time.Second)); err != nil {
t.Fatalf("AuthorizeObservation(owner): %v", err) t.Fatalf("AuthorizeObservation(owner): %v", err)
} }
if err := broker.CompleteObservation(context.Background(), "checker-a", observation, now.Add(time.Second)); err != nil { if err := broker.CompleteObservation(context.Background(), "checker-a", claimed[0].LeaseToken, observation, now.Add(time.Second)); err != nil {
t.Fatalf("CompleteObservation(): %v", err) t.Fatalf("CompleteObservation(): %v", err)
} }
if err := broker.AuthorizeObservation(context.Background(), "checker-a", observation, now.Add(2*time.Second)); err != nil { if err := broker.AuthorizeObservation(context.Background(), "checker-a", claimed[0].LeaseToken, observation, now.Add(2*time.Second)); err != nil {
t.Fatalf("AuthorizeObservation(replay): %v", err) t.Fatalf("AuthorizeObservation(replay): %v", err)
} }
} }
@ -82,7 +82,7 @@ func TestMemoryTaskBrokerReclaimsExpiredLeaseWithoutExceedingCheckerLimit(t *tes
} }
now = now.Add(2 * time.Second) now = now.Add(2 * time.Second)
reclaimed, err := broker.Claim(context.Background(), TaskClaim{CheckerID: "checker-b", InstanceID: "instance-b", MaxInFlight: 1, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}}) reclaimed, err := broker.Claim(context.Background(), TaskClaim{CheckerID: "checker-b", InstanceID: "instance-b", MaxInFlight: 1, SupportedLevels: []healthDomain.Level{healthDomain.LevelBasic}})
if err != nil || len(reclaimed) != 1 || reclaimed[0].TaskID != first[0].TaskID { if err != nil || len(reclaimed) != 1 || reclaimed[0].TaskID != first[0].TaskID || reclaimed[0].LeaseToken == first[0].LeaseToken {
t.Fatalf("reclaimed Claim() = (%+v, %v)", reclaimed, err) t.Fatalf("reclaimed Claim() = (%+v, %v)", reclaimed, err)
} }
} }