aiaa-notification-server/internal/subscriber/handle.go

package subscriber

import (
	"context"
	"errors"
	"log/slog"
	"strconv"

	"aiaa-notification-service/internal/model"
	"aiaa-notification-service/internal/notify"
	"aiaa-notification-service/internal/subscriber/tradesignal"
)

const retryHeader = "x-retry-count"

type Disposition int

const (
	DispositionAck Disposition = iota
	DispositionRetry
	DispositionDLQ
)

type SourceLookup func(ctx context.Context, name string) (*model.Source, error)

type ProcessFunc func(ctx context.Context, req notify.Request) (notify.Result, error)

type HandleInput struct {
	Body       []byte
	Headers    map[string]any
	Name       string
	Queue      string
	SourceName string
	MaxRetry   int
	Deduper    Deduper
}

func DecideRetry(retryCount, maxRetry int) Disposition {
	if retryCount+1 > maxRetry {
		return DispositionDLQ
	}
	return DispositionRetry
}

func RetryCount(headers map[string]any) int {
	if headers == nil {
		return 0
	}
	v, ok := headers[retryHeader]
	if !ok {
		return 0
	}
	switch n := v.(type) {
	case int:
		return n
	case int32:
		return int(n)
	case int64:
		return int(n)
	case float64:
		return int(n)
	case string:
		i, _ := strconv.Atoi(n)
		return i
	default:
		return 0
	}
}

type MessageConverter interface {
	Convert(body []byte) (event string, data map[string]interface{}, err error)
}

func errText(err error) string {
	if err == nil {
		return ""
	}
	return err.Error()
}

func HandleMessage(ctx context.Context, in HandleInput, conv MessageConverter, lookup SourceLookup, process ProcessFunc) Disposition {
	raw := string(in.Body)
	event, data, err := conv.Convert(in.Body)
	if err != nil {
		hash := MessageHash(in.Body)
		if errors.Is(err, tradesignal.ErrPositionStore) {
			slog.Warn("position store failed, retry", "hash", hash, "raw", raw, "error", errText(err))
			return DecideRetry(RetryCount(in.Headers), in.MaxRetry)
		}
		slog.Warn("invalid signal, ack", "hash", hash, "raw", raw, "error", errText(err))
		return DispositionAck
	}

	owned := false
	hash := SignalHash(in.SourceName, data)
	if in.Deduper != nil {
		ok, err := in.Deduper.Claim(ctx, hash)
		if err != nil {
			slog.Warn("dedup claim failed, processing anyway", "hash", hash, "error", errText(err))
		} else if !ok {
			slog.Info("duplicate message, ack", "hash", hash)
			return DispositionAck
		} else {
			owned = true
		}
	}

	slog.Info("mq message", "name", in.Name, "queue", in.Queue, "hash", hash, "raw", raw)

	src, err := lookup(ctx, in.SourceName)
	if err != nil || src == nil || src.Status != 1 {
		slog.Warn("source unavailable, ack", "source", in.SourceName, "hash", hash, "raw", raw, "error", errText(err))
		return DispositionAck
	}

	res, err := process(ctx, notify.Request{Source: src, Event: event, Data: data})
	if err == nil {
		if !res.Matched {
			slog.Info("no matching rule", "source", src.Name, "event", event, "hash", hash, "raw", raw)
		} else if res.Filtered {
			slog.Info("rule filtered", "source", src.Name, "event", event, "reason", res.Reason, "hash", hash, "raw", raw)
		} else {
			slog.Info("mq message accepted", "source", src.Name, "event", event, "channels", res.Channels, "hash", hash, "raw", raw)
		}
		return DispositionAck
	}
	if errors.Is(err, notify.ErrUnprocessable) {
		slog.Warn("unprocessable notify, ack", "source", src.Name, "event", event, "hash", hash, "raw", raw, "error", errText(err))
		return DispositionAck
	}

	disp := DecideRetry(RetryCount(in.Headers), in.MaxRetry)
	if owned && in.Deduper != nil {
		if relErr := in.Deduper.Release(ctx, hash); relErr != nil {
			slog.Warn("dedup release failed", "hash", hash, "error", errText(relErr))
		}
	}
	return disp
}