298 lines
10 KiB
Go
298 lines
10 KiB
Go
package redisactivity
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"proxy-pool/internal/domain/activitypool"
|
|
extractionDomain "proxy-pool/internal/domain/extraction"
|
|
healthDomain "proxy-pool/internal/domain/health"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
"proxy-pool/internal/platform/credentials"
|
|
)
|
|
|
|
const (
|
|
maxUpsertScriptBatch = 256
|
|
transientCredentialReleaseTimeout = time.Second
|
|
)
|
|
|
|
type upsertCandidate struct {
|
|
ProxyID string `json:"proxyId"`
|
|
UniqueDigest string `json:"uniqueDigest"`
|
|
Upstream string `json:"upstream"`
|
|
Record string `json:"record"`
|
|
}
|
|
|
|
var _ activitypool.Upserter = (*Adapter)(nil)
|
|
|
|
func (a *Adapter) UpsertFetched(ctx context.Context, upstreamID string, batch activitypool.FetchedBatch) (activitypool.UpsertResult, error) {
|
|
var result activitypool.UpsertResult
|
|
if a != nil {
|
|
defer a.releaseTransientCredentials(batch.Proxies)
|
|
}
|
|
if ctx == nil {
|
|
return result, activitypool.ErrInvalidBatch
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return result, err
|
|
}
|
|
if a == nil || upstreamID == "" || batch.ObservedAt.IsZero() || batch.ConfiguredTTL < 0 ||
|
|
batch.AllocationSafetyMargin < 0 || batch.MaxSize <= 0 ||
|
|
(batch.ConfiguredTTL > 0 && batch.AllocationSafetyMargin >= batch.ConfiguredTTL) {
|
|
return result, activitypool.ErrInvalidBatch
|
|
}
|
|
for _, candidate := range batch.Proxies {
|
|
if candidate.SourceUpstream != "" && candidate.SourceUpstream != upstreamID {
|
|
return result, activitypool.ErrInvalidBatch
|
|
}
|
|
}
|
|
|
|
prepared := make([]upsertCandidate, 0, len(batch.Proxies))
|
|
seenIDs := make(map[string]string, len(batch.Proxies))
|
|
for _, candidate := range batch.Proxies {
|
|
if err := ctx.Err(); err != nil {
|
|
return activitypool.UpsertResult{}, err
|
|
}
|
|
item, accepted, err := a.prepareUpsertCandidate(ctx, upstreamID, batch, candidate)
|
|
if err != nil {
|
|
return activitypool.UpsertResult{}, err
|
|
}
|
|
if !accepted {
|
|
result.Dropped++
|
|
continue
|
|
}
|
|
if digest, exists := seenIDs[item.ProxyID]; exists && digest != item.UniqueDigest {
|
|
return activitypool.UpsertResult{}, activitypool.ErrInvalidBatch
|
|
}
|
|
seenIDs[item.ProxyID] = item.UniqueDigest
|
|
prepared = append(prepared, item)
|
|
}
|
|
result.Accepted = len(prepared)
|
|
|
|
for start := 0; start < len(prepared); start += maxUpsertScriptBatch {
|
|
end := min(start+maxUpsertScriptBatch, len(prepared))
|
|
reply, err := a.upsertChunk(ctx, batch.ObservedAt, batch.MaxSize, prepared[start:end])
|
|
if err != nil {
|
|
return activitypool.UpsertResult{}, err
|
|
}
|
|
if reply.Accepted != end-start || reply.Inserted < 0 || reply.Refreshed < 0 || reply.Dropped < 0 ||
|
|
reply.Inserted+reply.Refreshed+reply.Dropped != reply.Accepted {
|
|
return activitypool.UpsertResult{}, invalidScriptReply("invalid upsert counters")
|
|
}
|
|
result.Inserted += reply.Inserted
|
|
result.Refreshed += reply.Refreshed
|
|
result.Dropped += reply.Dropped
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (a *Adapter) releaseTransientCredentials(proxies []proxyDomain.Proxy) {
|
|
if a.credentialReleaser == nil {
|
|
return
|
|
}
|
|
releaseCtx, cancel := context.WithTimeout(context.Background(), transientCredentialReleaseTimeout)
|
|
defer cancel()
|
|
for _, candidate := range proxies {
|
|
if releaseCtx.Err() != nil {
|
|
return
|
|
}
|
|
if candidate.SecretRef == "" || candidate.CredentialVersion == "" {
|
|
continue
|
|
}
|
|
_ = a.credentialReleaser.Release(releaseCtx, credentials.Reference{
|
|
SecretRef: candidate.SecretRef,
|
|
CredentialVersion: candidate.CredentialVersion,
|
|
})
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) prepareUpsertCandidate(
|
|
ctx context.Context,
|
|
upstreamID string,
|
|
batch activitypool.FetchedBatch,
|
|
candidate proxyDomain.Proxy,
|
|
) (upsertCandidate, bool, error) {
|
|
if !validCandidateIdentity(candidate) {
|
|
return upsertCandidate{}, false, nil
|
|
}
|
|
expiresAt := proxyDomain.EffectiveExpiry(batch.ObservedAt, candidate.ExpiresAt, 0, batch.ConfiguredTTL)
|
|
if expiresAt == nil {
|
|
return upsertCandidate{}, false, nil
|
|
}
|
|
usableUntil := expiresAt.Add(-batch.AllocationSafetyMargin)
|
|
if !usableUntil.After(batch.ObservedAt) || usableUntil.UnixMilli() <= batch.ObservedAt.UnixMilli() {
|
|
return upsertCandidate{}, false, nil
|
|
}
|
|
if candidate.MaxConcurrency < 0 || (candidate.State != "" && !validProxyState(string(candidate.State))) {
|
|
return upsertCandidate{}, false, activitypool.ErrInvalidBatch
|
|
}
|
|
if (candidate.SecretRef == "") != (candidate.CredentialVersion == "") {
|
|
return upsertCandidate{}, false, activitypool.ErrInvalidBatch
|
|
}
|
|
|
|
password := ""
|
|
if candidate.SecretRef != "" {
|
|
value, err := a.credentials.Resolve(ctx, credentials.Reference{
|
|
SecretRef: candidate.SecretRef, CredentialVersion: candidate.CredentialVersion,
|
|
})
|
|
if err != nil {
|
|
return upsertCandidate{}, false, fmt.Errorf("resolve proxy credential: %w", err)
|
|
}
|
|
candidate.Username = value.Username
|
|
password = value.Password
|
|
}
|
|
candidate.SourceUpstream = upstreamID
|
|
candidate.ExpiresAt = expiresAt
|
|
candidate.UsableUntil = &usableUntil
|
|
if candidate.CreatedAt.IsZero() {
|
|
candidate.CreatedAt = batch.ObservedAt.UTC()
|
|
}
|
|
if candidate.State == "" {
|
|
candidate.State = proxyDomain.StateFetched
|
|
}
|
|
uniqueKey := candidate.UniqueKey()
|
|
if candidate.ID == "" {
|
|
candidate.ID = stableAdapterProxyID(uniqueKey)
|
|
}
|
|
record := proxyRecord{
|
|
Version: recordVersion, ID: candidate.ID, Scheme: string(candidate.Scheme), Host: candidate.Host,
|
|
Port: int64(candidate.Port), Username: candidate.Username, Password: password,
|
|
CredentialVersion: candidate.CredentialVersion, SourceUpstream: upstreamID,
|
|
CreatedAtMS: candidate.CreatedAt.UnixMilli(), ExpiresAtMS: expiresAt.UnixMilli(),
|
|
UsableUntilMS: usableUntil.UnixMilli(), LatencyNS: int64(candidate.Latency),
|
|
MaxConcurrency: candidate.MaxConcurrency, State: string(candidate.State),
|
|
Tags: cloneTags(candidate.Tags), OwnerIndexKey: a.keys.owned(upstreamID),
|
|
IndexKeys: a.availableIndexKeys(candidate),
|
|
}
|
|
if candidate.LastCheckedAt != nil {
|
|
record.LastCheckedAtMS = candidate.LastCheckedAt.UnixMilli()
|
|
}
|
|
if candidate.LastSuccessAt != nil {
|
|
record.LastSuccessAtMS = candidate.LastSuccessAt.UnixMilli()
|
|
}
|
|
encoded, err := encodeProxyRecord(record)
|
|
if err != nil {
|
|
if errors.Is(err, ErrInvalidRecord) {
|
|
return upsertCandidate{}, false, activitypool.ErrInvalidBatch
|
|
}
|
|
return upsertCandidate{}, false, err
|
|
}
|
|
return upsertCandidate{
|
|
ProxyID: candidate.ID, UniqueDigest: digestToken(uniqueKey), Upstream: upstreamID, Record: encoded,
|
|
}, true, nil
|
|
}
|
|
|
|
func (a *Adapter) upsertChunk(
|
|
ctx context.Context,
|
|
observedAt time.Time,
|
|
maxSize int,
|
|
candidates []upsertCandidate,
|
|
) (upsertScriptReply, error) {
|
|
operationID, err := newOperationID()
|
|
if err != nil {
|
|
return upsertScriptReply{}, err
|
|
}
|
|
payload, err := json.Marshal(candidates)
|
|
if err != nil {
|
|
return upsertScriptReply{}, fmt.Errorf("encode Redis upsert candidates: %w", err)
|
|
}
|
|
result, err := runScript(ctx, a.client, upsertScript, []string{
|
|
a.keys.records, a.keys.unique, a.keys.idkeys, a.keys.expiry, a.keys.available,
|
|
a.keys.inventory, a.keys.stateInventory, a.keys.owners, a.keys.ownerExpiry,
|
|
a.keys.operation(operationID),
|
|
}, observedAt.UnixMilli(), a.options.CleanupLimit, maxSize, operationTTLMillis(a.options.OperationTTL), string(payload))
|
|
if err != nil {
|
|
return upsertScriptReply{}, err
|
|
}
|
|
var reply upsertScriptReply
|
|
if err := decodeScriptResult(result, &reply); err != nil {
|
|
return upsertScriptReply{}, err
|
|
}
|
|
switch reply.Status {
|
|
case scriptOK:
|
|
return reply, nil
|
|
case scriptInvalid:
|
|
return upsertScriptReply{}, activitypool.ErrInvalidBatch
|
|
default:
|
|
return upsertScriptReply{}, invalidScriptReply("unexpected upsert status")
|
|
}
|
|
}
|
|
|
|
func (a *Adapter) availableIndexKeys(candidate proxyDomain.Proxy) []string {
|
|
keys := []string{a.keys.protocol(string(candidate.Scheme))}
|
|
if region := candidate.Tags["region"]; region != "" {
|
|
keys = append(keys, a.keys.region(region))
|
|
}
|
|
if carrier := candidate.Tags["carrier"]; carrier != "" {
|
|
keys = append(keys, a.keys.carrier(carrier))
|
|
}
|
|
return append(keys, a.keys.upstream(candidate.SourceUpstream))
|
|
}
|
|
|
|
func proxyRecordEntry(record proxyRecord) activitypool.Entry {
|
|
createdAt := time.UnixMilli(record.CreatedAtMS).UTC()
|
|
expiresAt := time.UnixMilli(record.ExpiresAtMS).UTC()
|
|
usableUntil := time.UnixMilli(record.UsableUntilMS).UTC()
|
|
proxy := proxyDomain.Proxy{
|
|
ID: record.ID, Scheme: proxyDomain.Scheme(record.Scheme), Host: record.Host, Port: uint16(record.Port),
|
|
Username: record.Username, CredentialVersion: record.CredentialVersion,
|
|
SourceUpstream: record.SourceUpstream, CreatedAt: createdAt, ExpiresAt: &expiresAt,
|
|
UsableUntil: &usableUntil, Latency: time.Duration(record.LatencyNS),
|
|
MaxConcurrency: record.MaxConcurrency, State: proxyDomain.State(record.State), Tags: cloneTags(record.Tags),
|
|
}
|
|
if record.LastCheckedAtMS > 0 {
|
|
value := time.UnixMilli(record.LastCheckedAtMS).UTC()
|
|
proxy.LastCheckedAt = &value
|
|
}
|
|
if record.LastSuccessAtMS > 0 {
|
|
value := time.UnixMilli(record.LastSuccessAtMS).UTC()
|
|
proxy.LastSuccessAt = &value
|
|
}
|
|
global := healthDomain.GlobalState{State: proxyDomain.State(record.State)}
|
|
if record.LastHealthTaskID != "" {
|
|
digest, err := hex.DecodeString(record.LastHealthDigest)
|
|
if err == nil && len(digest) == sha256.Size && record.LastHealthObservedAtMS > 0 {
|
|
copy(global.LastObservationDigest[:], digest)
|
|
global.LastTaskID = record.LastHealthTaskID
|
|
global.LastObservedAt = time.UnixMilli(record.LastHealthObservedAtMS).UTC()
|
|
global.ConsecutiveFailures = record.ConsecutiveFailures
|
|
}
|
|
}
|
|
return activitypool.Entry{
|
|
Proxy: proxy, UsableUntil: usableUntil,
|
|
OwnerWorkerID: record.OwnerWorkerID, State: proxyDomain.State(record.State), GlobalHealth: global,
|
|
}
|
|
}
|
|
|
|
func validCandidateIdentity(candidate proxyDomain.Proxy) bool {
|
|
if candidate.Host == "" || candidate.Port == 0 {
|
|
return false
|
|
}
|
|
return validScheme(string(candidate.Scheme))
|
|
}
|
|
|
|
func stableAdapterProxyID(uniqueKey string) string {
|
|
digest := sha256.Sum256([]byte(uniqueKey))
|
|
return "px_" + hex.EncodeToString(digest[:12])
|
|
}
|
|
|
|
func cloneTags(tags map[string]string) map[string]string {
|
|
if tags == nil {
|
|
return nil
|
|
}
|
|
cloned := make(map[string]string, len(tags))
|
|
for key, value := range tags {
|
|
cloned[key] = value
|
|
}
|
|
return cloned
|
|
}
|
|
|
|
func invalidScriptReply(reason string) error {
|
|
return errors.Join(extractionDomain.ErrStoreUnavailable, errors.New(reason))
|
|
}
|