First commit Go counter service
This commit is contained in:
@@ -0,0 +1,345 @@
|
||||
// Package counting turns one ingested recipient into the counter increments
|
||||
// it causes, and executes them against the store.
|
||||
//
|
||||
// Rules (see plans/2026-10-13-migadu-counters.md):
|
||||
//
|
||||
// - /out (authenticated submission): the sender's A/D/O outgoing counters
|
||||
// +1 when the sender's domain is hosted here, and the recipient's A/D/O
|
||||
// incoming counters +1 when the recipient's domain is hosted here.
|
||||
// - /in (inbound from outside): the recipient's A/D/O incoming counters +1
|
||||
// when hosted here; the sender is never counted; junk is skipped unless
|
||||
// configured otherwise.
|
||||
// - Mail to or from an exempt address (support@, contact@) never counts.
|
||||
// - A domain that is not in the domains CDB is not counted.
|
||||
//
|
||||
// The Planner is pure (no I/O besides the domain lookup) so the rules are
|
||||
// testable without Redis; Counter runs a Plan's increments in one pipeline.
|
||||
package counting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"strconv"
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"migadu_counters/internal/domains"
|
||||
"migadu_counters/internal/store"
|
||||
)
|
||||
|
||||
// Routes and skip reasons as reported in responses and the transaction log.
|
||||
const (
|
||||
RouteIn = "in"
|
||||
RouteOut = "out"
|
||||
|
||||
SkipExempt = "exempt"
|
||||
SkipJunk = "junk"
|
||||
SkipUnknownDomain = "unknown_domain"
|
||||
)
|
||||
|
||||
// Lookuper resolves a domain to its CDB record; *domains.DB implements it.
|
||||
type Lookuper interface {
|
||||
Lookup(domain string) (*domains.Domain, error)
|
||||
}
|
||||
|
||||
// LookupFunc adapts a function (e.g. one reading an atomic.Pointer to the
|
||||
// current DB) to Lookuper.
|
||||
type LookupFunc func(domain string) (*domains.Domain, error)
|
||||
|
||||
// Lookup implements Lookuper.
|
||||
func (f LookupFunc) Lookup(domain string) (*domains.Domain, error) { return f(domain) }
|
||||
|
||||
// Planner applies the counting rules.
|
||||
type Planner struct {
|
||||
Lookup Lookuper
|
||||
Exempt map[string]bool // lowercased addresses that never count
|
||||
CountJunk bool // count /in requests flagged junk
|
||||
Now func() time.Time // nil = time.Now; the UTC day of Now() is the counter day
|
||||
}
|
||||
|
||||
// Job is one validated request handed to the queue: the fields the counting
|
||||
// rules need and nothing else. A side that failed validation is empty, with
|
||||
// the reason in SenderUnusable / RecipientUnusable for the log.
|
||||
type Job struct {
|
||||
Route string // RouteIn or RouteOut
|
||||
TraceID string
|
||||
Received time.Time // when the request arrived (queue wait metric)
|
||||
|
||||
AuthUser string // /out: normalised SMTP AUTH user, "" when not usable
|
||||
LoginSuffix string // /out: "SUPPORTTOKEN" of a user@dom@SUPPORTTOKEN login
|
||||
MailFrom string // normalised, or the sanitised input when not an address
|
||||
// MailFromUsable tells whether MailFrom is a normalised address (and may
|
||||
// match the exempt list); an unusable value is never exempt.
|
||||
MailFromUsable bool
|
||||
Recipient string // normalised, "" when not usable
|
||||
Junk bool // /in only
|
||||
|
||||
SenderUnusable Unusable // /out: auth_user not usable (Value == "" when usable)
|
||||
RecipientUnusable Unusable // recipient not usable
|
||||
}
|
||||
|
||||
// Unusable records an address side that could not be counted.
|
||||
type Unusable struct {
|
||||
Role string // "sender" or "recipient"
|
||||
Value string // sanitised input, "" when the field was missing
|
||||
Reason string // "missing" or "invalid e-mail address"
|
||||
}
|
||||
|
||||
// Set reports whether a side was recorded as unusable.
|
||||
func (u Unusable) Set() bool { return u.Role != "" }
|
||||
|
||||
// Lookup records one domain resolution for the transaction log.
|
||||
type Lookup struct {
|
||||
Role string // "sender" or "recipient"
|
||||
Address string
|
||||
DomainName string
|
||||
Domain *domains.Domain // nil when not found
|
||||
Err error // domains.ErrNotFound when not hosted here
|
||||
}
|
||||
|
||||
// Increment is one counter the request bumps.
|
||||
type Increment struct {
|
||||
Level string // A, D or O
|
||||
ID string // address, domain name or organisation id
|
||||
Flow string // i or o
|
||||
Key string // Redis key for the day
|
||||
By string // address that caused it (sender for o, recipient for i); $ADDRESS in notifications
|
||||
Limit int64 // configured limit, 0 when unlimited
|
||||
Limited bool // false when no limit is configured for this level/flow
|
||||
Domain *domains.Domain
|
||||
}
|
||||
|
||||
// Plan is the result of applying the rules to one request.
|
||||
type Plan struct {
|
||||
Route string
|
||||
TraceID string
|
||||
Date string // UTC day used in the keys
|
||||
Skipped string // "" or one of the Skip* reasons
|
||||
AuthUser string
|
||||
LoginSuffix string
|
||||
MailFrom string
|
||||
Recipient string
|
||||
Junk bool // /in only: the message was flagged junk by mizu
|
||||
|
||||
Unusable []Unusable // sides that were not counted because they failed validation
|
||||
Lookups []Lookup
|
||||
Increments []Increment
|
||||
}
|
||||
|
||||
func (p *Planner) now() time.Time {
|
||||
if p.Now != nil {
|
||||
return p.Now()
|
||||
}
|
||||
return time.Now()
|
||||
}
|
||||
|
||||
func (p *Planner) isExempt(addr string) bool { return addr != "" && p.Exempt[addr] }
|
||||
|
||||
func normalize(addr string) string { return strings.ToLower(strings.TrimSpace(addr)) }
|
||||
|
||||
// splitAddress returns local part and domain; ok is false without an "@".
|
||||
func splitAddress(addr string) (local, domain string, ok bool) {
|
||||
i := strings.LastIndex(addr, "@")
|
||||
if i <= 0 || i == len(addr)-1 {
|
||||
return "", "", false
|
||||
}
|
||||
return addr[:i], addr[i+1:], true
|
||||
}
|
||||
|
||||
// PlanOut applies the /out rules for one authenticated message recipient
|
||||
// whose fields all validated.
|
||||
func (p *Planner) PlanOut(traceID, authUser, mailFrom, recipient string) (*Plan, error) {
|
||||
return p.Plan(Job{Route: RouteOut, TraceID: traceID, AuthUser: authUser, MailFrom: mailFrom, MailFromUsable: mailFrom != "", Recipient: recipient})
|
||||
}
|
||||
|
||||
// PlanIn applies the /in rules for one inbound recipient. mailFrom may be
|
||||
// empty (bounce sender <>).
|
||||
func (p *Planner) PlanIn(traceID, mailFrom, recipient string, junk bool) (*Plan, error) {
|
||||
return p.Plan(Job{Route: RouteIn, TraceID: traceID, MailFrom: mailFrom, MailFromUsable: mailFrom != "", Recipient: recipient, Junk: junk})
|
||||
}
|
||||
|
||||
// Plan applies the counting rules to a validated Job. Unusable sides are
|
||||
// recorded for the log and skipped; the usable side is counted.
|
||||
func (p *Planner) Plan(job Job) (*Plan, error) {
|
||||
authUser, mailFrom, recipient := normalize(job.AuthUser), job.MailFrom, normalize(job.Recipient)
|
||||
if job.MailFromUsable {
|
||||
mailFrom = normalize(mailFrom)
|
||||
}
|
||||
plan := &Plan{Route: job.Route, TraceID: job.TraceID, Date: store.DayKey(p.now()),
|
||||
AuthUser: authUser, LoginSuffix: job.LoginSuffix, MailFrom: mailFrom, Recipient: recipient, Junk: job.Junk}
|
||||
if job.SenderUnusable.Set() {
|
||||
plan.Unusable = append(plan.Unusable, job.SenderUnusable)
|
||||
}
|
||||
if job.RecipientUnusable.Set() {
|
||||
plan.Unusable = append(plan.Unusable, job.RecipientUnusable)
|
||||
}
|
||||
if p.isExempt(domains.Mailbox(authUser)) || (job.MailFromUsable && p.isExempt(domains.Mailbox(mailFrom))) || p.isExempt(domains.Mailbox(recipient)) {
|
||||
plan.Skipped = SkipExempt
|
||||
return plan, nil
|
||||
}
|
||||
if job.Route == RouteIn && job.Junk && !p.CountJunk {
|
||||
plan.Skipped = SkipJunk
|
||||
return plan, nil
|
||||
}
|
||||
if job.Route == RouteOut && authUser != "" {
|
||||
if err := p.add(plan, "sender", authUser, domains.FlowOut); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if recipient != "" {
|
||||
if err := p.add(plan, "recipient", recipient, domains.FlowIn); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if len(plan.Increments) == 0 {
|
||||
plan.Skipped = SkipUnknownDomain
|
||||
}
|
||||
return plan, nil
|
||||
}
|
||||
|
||||
// add looks up the address's domain and, when hosted here, appends its
|
||||
// A, D and O increments for flow. The address is counted as its mailbox
|
||||
// (plus-tag stripped). A not-found or malformed domain row is recorded and
|
||||
// skipped; any other lookup error (unreadable CDB) aborts the plan.
|
||||
func (p *Planner) add(plan *Plan, role, address, flow string) error {
|
||||
address = domains.Mailbox(address)
|
||||
local, domain, ok := splitAddress(address)
|
||||
lk := Lookup{Role: role, Address: address, DomainName: domain}
|
||||
if !ok {
|
||||
lk.Err = domains.ErrNotFound
|
||||
plan.Lookups = append(plan.Lookups, lk)
|
||||
return nil
|
||||
}
|
||||
d, err := p.Lookup.Lookup(domain)
|
||||
if err != nil {
|
||||
lk.Err = err
|
||||
plan.Lookups = append(plan.Lookups, lk)
|
||||
if errors.Is(err, domains.ErrNotFound) || errors.Is(err, domains.ErrMalformed) {
|
||||
return nil
|
||||
}
|
||||
return fmt.Errorf("lookup %s domain %s: %w", role, domain, err)
|
||||
}
|
||||
lk.Domain = d
|
||||
plan.Lookups = append(plan.Lookups, lk)
|
||||
|
||||
levels := []struct{ level, id string }{
|
||||
{domains.AddressLevel, address},
|
||||
{domains.DomainLevel, d.Name},
|
||||
{domains.OrganizationLevel, strconv.FormatInt(d.OrgID, 10)},
|
||||
}
|
||||
for _, l := range levels {
|
||||
limit, limited := d.Limits.Limit(l.level, flow, local)
|
||||
plan.Increments = append(plan.Increments, Increment{
|
||||
Level: l.level, ID: l.id, Flow: flow,
|
||||
Key: store.Key(l.level, l.id, plan.Date, flow),
|
||||
By: address, Limit: limit, Limited: limited, Domain: d,
|
||||
})
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Crossed reports whether an increment of amount that produced newValue is
|
||||
// the one that reached limit: old < limit <= new. Because INCRBY is atomic
|
||||
// exactly one increment can cross, so this is the once-per-crossing trigger
|
||||
// for notifications; a limit of 0 (or unlimited) never crosses.
|
||||
func Crossed(newValue, amount, limit int64) bool {
|
||||
return limit > 0 && newValue-amount < limit && limit <= newValue
|
||||
}
|
||||
|
||||
// Result is an executed Increment.
|
||||
type Result struct {
|
||||
Increment
|
||||
Count int64 // value after the increment (from EXEC, or re-read after a lost reply)
|
||||
Reached bool // this increment crossed the limit
|
||||
Err error // this key's own INCR failed (e.g. it holds a non-number); Count is 0
|
||||
}
|
||||
|
||||
// Outcome is an executed Plan.
|
||||
type Outcome struct {
|
||||
*Plan
|
||||
Results []Result
|
||||
// Unknown is set when the EXEC reply was lost after the write; Counts
|
||||
// were then re-read from the master and Reached means count >= limit.
|
||||
Unknown error
|
||||
}
|
||||
|
||||
// Counter plans and executes.
|
||||
type Counter struct {
|
||||
Planner *Planner
|
||||
Store *store.Store
|
||||
ReReadTimeout time.Duration // budget for the master re-read after a lost reply (default 2s)
|
||||
}
|
||||
|
||||
// CountOut plans and executes an /out request.
|
||||
func (c *Counter) CountOut(ctx context.Context, traceID, authUser, mailFrom, recipient string) (*Outcome, error) {
|
||||
plan, err := c.Planner.PlanOut(traceID, authUser, mailFrom, recipient)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c.Execute(ctx, plan)
|
||||
}
|
||||
|
||||
// CountIn plans and executes an /in request.
|
||||
func (c *Counter) CountIn(ctx context.Context, traceID, mailFrom, recipient string, junk bool) (*Outcome, error) {
|
||||
plan, err := c.Planner.PlanIn(traceID, mailFrom, recipient, junk)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return c.Execute(ctx, plan)
|
||||
}
|
||||
|
||||
// Execute runs all increments of plan in one pipeline. A skipped plan does
|
||||
// not touch Redis.
|
||||
//
|
||||
// Failure handling (review #4): a key whose own INCR failed is reported in
|
||||
// its Result.Err while the other keys are counted normally. When the whole
|
||||
// reply was lost (timeout, reset) after the write may have been applied, the
|
||||
// keys are re-read once from the master and the outcome is flagged Unknown;
|
||||
// crossing detection then uses count >= limit so a crossing that happened in
|
||||
// the lost transaction is still notified. A definite failure (connection
|
||||
// refused, cancelled) is returned as an error: nothing was counted.
|
||||
func (c *Counter) Execute(ctx context.Context, plan *Plan) (*Outcome, error) {
|
||||
out := &Outcome{Plan: plan}
|
||||
if len(plan.Increments) == 0 {
|
||||
return out, nil
|
||||
}
|
||||
keys := make([]string, len(plan.Increments))
|
||||
for i, inc := range plan.Increments {
|
||||
keys[i] = inc.Key
|
||||
}
|
||||
out.Results = make([]Result, len(plan.Increments))
|
||||
results, err := c.Store.Incr(ctx, keys)
|
||||
if err == nil {
|
||||
for i, inc := range plan.Increments {
|
||||
r := Result{Increment: inc, Err: results[i].Err}
|
||||
if r.Err == nil {
|
||||
r.Count = results[i].Count
|
||||
r.Reached = inc.Limited && Crossed(r.Count, 1, inc.Limit)
|
||||
}
|
||||
out.Results[i] = r
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
if !store.OutcomeUnknown(err) {
|
||||
return nil, err
|
||||
}
|
||||
rctx, cancel := context.WithTimeout(context.WithoutCancel(ctx), c.reReadTimeout())
|
||||
defer cancel()
|
||||
counts, rerr := c.Store.GetMaster(rctx, keys)
|
||||
if rerr != nil {
|
||||
return nil, fmt.Errorf("%w (outcome unknown, re-read failed: %v)", err, rerr)
|
||||
}
|
||||
out.Unknown = err
|
||||
for i, inc := range plan.Increments {
|
||||
out.Results[i] = Result{Increment: inc, Count: counts[i], Reached: inc.Limited && counts[i] >= inc.Limit}
|
||||
}
|
||||
return out, nil
|
||||
}
|
||||
|
||||
func (c *Counter) reReadTimeout() time.Duration {
|
||||
if c.ReReadTimeout > 0 {
|
||||
return c.ReReadTimeout
|
||||
}
|
||||
return 2 * time.Second
|
||||
}
|
||||
@@ -0,0 +1,677 @@
|
||||
package counting
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/alicebob/miniredis/v2"
|
||||
"github.com/redis/go-redis/v9"
|
||||
|
||||
"migadu_counters/internal/domains"
|
||||
"migadu_counters/internal/store"
|
||||
)
|
||||
|
||||
// table is a fake domains lookup: domain -> record. Unknown domains return
|
||||
// domains.ErrNotFound; the special domain "broken.tld" returns a read error.
|
||||
type table map[string]*domains.Domain
|
||||
|
||||
func (t table) Lookup(domain string) (*domains.Domain, error) {
|
||||
if domain == "broken.tld" {
|
||||
return nil, errors.New("cdb read error")
|
||||
}
|
||||
if d, ok := t[domain]; ok {
|
||||
return d, nil
|
||||
}
|
||||
return nil, domains.ErrNotFound
|
||||
}
|
||||
|
||||
var (
|
||||
fixedNow = time.Date(2026, 10, 13, 12, 0, 0, 0, time.UTC)
|
||||
day = "2026-10-13"
|
||||
exempt = map[string]bool{"support@migadu.com": true, "contact@migadu.com": true}
|
||||
)
|
||||
|
||||
func newTable() table {
|
||||
return table{
|
||||
"example.com": {Name: "example.com", OrgID: 7,
|
||||
Limits: domains.ParseLimits("O/i=10000,O/o=2200,D/i=10000,D/o=2200,A/i=500,A/o=100,noreply@/o=5")},
|
||||
"other.org": {Name: "other.org", OrgID: 9,
|
||||
Limits: domains.ParseLimits("O/i=200,O/o=20,D/i=200,D/o=20,A/i=200,A/o=20")},
|
||||
"nolimits.net": {Name: "nolimits.net", OrgID: 11, Limits: domains.Limits{}},
|
||||
}
|
||||
}
|
||||
|
||||
func newPlanner(countJunk bool) *Planner {
|
||||
return &Planner{Lookup: newTable(), Exempt: exempt, CountJunk: countJunk, Now: func() time.Time { return fixedNow }}
|
||||
}
|
||||
|
||||
type want struct {
|
||||
level, id, flow, by string
|
||||
limit int64
|
||||
limited bool
|
||||
}
|
||||
|
||||
func checkIncrements(t *testing.T, got []Increment, wants []want) {
|
||||
t.Helper()
|
||||
if len(got) != len(wants) {
|
||||
t.Fatalf("got %d increments, want %d:\n%s", len(got), len(wants), dump(got))
|
||||
}
|
||||
for i, w := range wants {
|
||||
g := got[i]
|
||||
wantKey := store.Key(w.level, w.id, day, w.flow)
|
||||
if g.Level != w.level || g.ID != w.id || g.Flow != w.flow || g.By != w.by || g.Limit != w.limit || g.Limited != w.limited || g.Key != wantKey {
|
||||
t.Errorf("increment %d = %+v,\n want level=%s id=%s flow=%s by=%s limit=%d limited=%v key=%s",
|
||||
i, g, w.level, w.id, w.flow, w.by, w.limit, w.limited, wantKey)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func dump(incs []Increment) string {
|
||||
s := ""
|
||||
for i, in := range incs {
|
||||
s += fmt.Sprintf(" %d: %+v\n", i, in)
|
||||
}
|
||||
return s
|
||||
}
|
||||
|
||||
// x@example.com -> y@example.com: sender outgoing on A/D/O, recipient incoming on A/D/O.
|
||||
func TestPlanOut_LocalRecipient(t *testing.T) {
|
||||
p := newPlanner(false)
|
||||
plan, err := p.PlanOut("tr1", "x@example.com", "x@example.com", "y@example.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != "" || plan.Route != "out" || plan.TraceID != "tr1" || plan.Date != day {
|
||||
t.Errorf("plan header = %+v", plan)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "x@example.com", "o", "x@example.com", 100, true},
|
||||
{"D", "example.com", "o", "x@example.com", 2200, true},
|
||||
{"O", "7", "o", "x@example.com", 2200, true},
|
||||
{"A", "y@example.com", "i", "y@example.com", 500, true},
|
||||
{"D", "example.com", "i", "y@example.com", 10000, true},
|
||||
{"O", "7", "i", "y@example.com", 10000, true},
|
||||
})
|
||||
if len(plan.Lookups) != 2 || plan.Lookups[0].Role != "sender" || plan.Lookups[0].Domain == nil ||
|
||||
plan.Lookups[1].Role != "recipient" || plan.Lookups[1].Domain == nil {
|
||||
t.Errorf("lookups = %+v", plan.Lookups)
|
||||
}
|
||||
}
|
||||
|
||||
// x@example.com -> s@gmail.com: only the sender side; the recipient lookup is recorded as not found.
|
||||
func TestPlanOut_ExternalRecipient(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanOut("tr", "x@example.com", "x@example.com", "s@gmail.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "x@example.com", "o", "x@example.com", 100, true},
|
||||
{"D", "example.com", "o", "x@example.com", 2200, true},
|
||||
{"O", "7", "o", "x@example.com", 2200, true},
|
||||
})
|
||||
if len(plan.Lookups) != 2 || plan.Lookups[1].Domain != nil || !errors.Is(plan.Lookups[1].Err, domains.ErrNotFound) {
|
||||
t.Errorf("recipient lookup should be recorded as not found: %+v", plan.Lookups)
|
||||
}
|
||||
}
|
||||
|
||||
// Recipient in another hosted org: its own domain row decides D/O ids and limits.
|
||||
func TestPlanOut_RecipientInOtherOrg(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanOut("tr", "x@example.com", "", "z@other.org")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
checkIncrements(t, plan.Increments[3:], []want{
|
||||
{"A", "z@other.org", "i", "z@other.org", 200, true},
|
||||
{"D", "other.org", "i", "z@other.org", 200, true},
|
||||
{"O", "9", "i", "z@other.org", 200, true},
|
||||
})
|
||||
}
|
||||
|
||||
func TestPlanOut_SenderDomainUnknown(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanOut("tr", "staff@unknown.tld", "", "y@example.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != "" {
|
||||
t.Errorf("skipped = %q, want counting of the local recipient", plan.Skipped)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "y@example.com", "i", "y@example.com", 500, true},
|
||||
{"D", "example.com", "i", "y@example.com", 10000, true},
|
||||
{"O", "7", "i", "y@example.com", 10000, true},
|
||||
})
|
||||
}
|
||||
|
||||
func TestPlanOut_BothUnknown(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanOut("tr", "a@unknown.tld", "", "b@gmail.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != SkipUnknownDomain || len(plan.Increments) != 0 {
|
||||
t.Errorf("plan = %+v", plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanOut_Exempt(t *testing.T) {
|
||||
cases := []struct{ authUser, mailFrom, recipient string }{
|
||||
{"support@migadu.com", "support@migadu.com", "y@example.com"},
|
||||
{"x@example.com", "x@example.com", "support@migadu.com"},
|
||||
{"x@example.com", "contact@migadu.com", "y@example.com"},
|
||||
{"x@example.com", "x@example.com", "Contact@Migadu.com"},
|
||||
}
|
||||
for _, c := range cases {
|
||||
plan, err := newPlanner(false).PlanOut("tr", c.authUser, c.mailFrom, c.recipient)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != SkipExempt || len(plan.Increments) != 0 || len(plan.Lookups) != 0 {
|
||||
t.Errorf("%+v: plan = %+v, want exempt with no lookups", c, plan)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanOut_AddressOverrideAndUnlimited(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanOut("tr", "noreply@example.com", "", "q@nolimits.net")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "noreply@example.com", "o", "noreply@example.com", 5, true}, // noreply@/o override
|
||||
{"D", "example.com", "o", "noreply@example.com", 2200, true},
|
||||
{"O", "7", "o", "noreply@example.com", 2200, true},
|
||||
{"A", "q@nolimits.net", "i", "q@nolimits.net", 0, false},
|
||||
{"D", "nolimits.net", "i", "q@nolimits.net", 0, false},
|
||||
{"O", "11", "i", "q@nolimits.net", 0, false},
|
||||
})
|
||||
}
|
||||
|
||||
func TestPlanOut_NormalizesAddresses(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanOut("tr", " X@Example.COM ", "", "Y@EXAMPLE.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Increments[0].ID != "x@example.com" || plan.Increments[3].ID != "y@example.com" {
|
||||
t.Errorf("ids not normalised: %s", dump(plan.Increments))
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanIn_LocalRecipient(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanIn("tr9", "s@gmail.com", "y@example.com", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Route != "in" || plan.TraceID != "tr9" || plan.Skipped != "" {
|
||||
t.Errorf("plan header = %+v", plan)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "y@example.com", "i", "y@example.com", 500, true},
|
||||
{"D", "example.com", "i", "y@example.com", 10000, true},
|
||||
{"O", "7", "i", "y@example.com", 10000, true},
|
||||
})
|
||||
if len(plan.Lookups) != 1 || plan.Lookups[0].Role != "recipient" {
|
||||
t.Errorf("lookups = %+v", plan.Lookups)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanIn_SenderNeverCounted(t *testing.T) {
|
||||
// A hosted sender on /in (e.g. mail relayed back in) still only counts the recipient.
|
||||
plan, err := newPlanner(false).PlanIn("tr", "x@example.com", "z@other.org", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
for _, in := range plan.Increments {
|
||||
if in.Flow != "i" || in.By != "z@other.org" {
|
||||
t.Errorf("unexpected increment %+v", in)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanIn_UnknownDomain(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanIn("tr", "s@gmail.com", "nobody@gmail.com", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != SkipUnknownDomain || len(plan.Increments) != 0 {
|
||||
t.Errorf("plan = %+v", plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanIn_Junk(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanIn("tr", "s@gmail.com", "y@example.com", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != SkipJunk || len(plan.Increments) != 0 || len(plan.Lookups) != 0 {
|
||||
t.Errorf("junk not skipped: %+v", plan)
|
||||
}
|
||||
plan, err = newPlanner(true).PlanIn("tr", "s@gmail.com", "y@example.com", true)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != "" || len(plan.Increments) != 3 {
|
||||
t.Errorf("count_junk=true should count: %+v", plan)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanIn_Exempt(t *testing.T) {
|
||||
for _, c := range []struct{ from, rcpt string }{{"support@migadu.com", "y@example.com"}, {"s@gmail.com", "contact@migadu.com"}} {
|
||||
plan, err := newPlanner(false).PlanIn("tr", c.from, c.rcpt, false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != SkipExempt || len(plan.Increments) != 0 {
|
||||
t.Errorf("%+v: plan = %+v", c, plan)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlanIn_EmptyMailFromIsBounce(t *testing.T) {
|
||||
plan, err := newPlanner(false).PlanIn("tr", "", "y@example.com", false)
|
||||
if err != nil || plan.Skipped != "" || len(plan.Increments) != 3 {
|
||||
t.Errorf("bounce (<>) must still count the recipient: %+v, %v", plan, err)
|
||||
}
|
||||
}
|
||||
|
||||
// A CDB read error (not "not found") must surface so the server answers 503.
|
||||
func TestPlan_LookupError(t *testing.T) {
|
||||
if _, err := newPlanner(false).PlanIn("tr", "s@gmail.com", "y@broken.tld", false); err == nil {
|
||||
t.Error("PlanIn: expected error")
|
||||
}
|
||||
if _, err := newPlanner(false).PlanOut("tr", "x@broken.tld", "", "y@example.com"); err == nil {
|
||||
t.Error("PlanOut: expected error")
|
||||
}
|
||||
}
|
||||
|
||||
// Against the real CDB: yagolivas.com is org 1049 with A/o=20, D/o=20, O/o=20.
|
||||
func TestPlanOut_RealCDB(t *testing.T) {
|
||||
db, err := domains.Open("../../test_files/rspamd_domains_db.cdb")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
p := &Planner{Lookup: db, Exempt: exempt, Now: func() time.Time { return fixedNow }}
|
||||
plan, err := p.PlanOut("tr", "someone@yagolivas.com", "", "ext@gmail.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "someone@yagolivas.com", "o", "someone@yagolivas.com", 20, true},
|
||||
{"D", "yagolivas.com", "o", "someone@yagolivas.com", 20, true},
|
||||
{"O", "1049", "o", "someone@yagolivas.com", 20, true},
|
||||
})
|
||||
}
|
||||
|
||||
func TestCrossed(t *testing.T) {
|
||||
tests := []struct {
|
||||
newValue, amount, limit int64
|
||||
want bool
|
||||
}{
|
||||
{1, 1, 1, true},
|
||||
{5, 1, 5, true},
|
||||
{4, 1, 5, false},
|
||||
{6, 1, 5, false}, // already over: no second notification
|
||||
{5, 3, 5, true}, // 2 -> 5 crosses
|
||||
{7, 3, 5, true}, // 4 -> 7 crosses
|
||||
{8, 3, 5, false}, // 5 -> 8 was already at the limit
|
||||
{1, 1, 0, false}, // 0 / unlimited never crosses
|
||||
{1, 1, -1, false},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
if got := Crossed(tt.newValue, tt.amount, tt.limit); got != tt.want {
|
||||
t.Errorf("Crossed(%d, %d, %d) = %v, want %v", tt.newValue, tt.amount, tt.limit, got, tt.want)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func newCounter(t *testing.T) (*Counter, *miniredis.Miniredis) {
|
||||
t.Helper()
|
||||
m := miniredis.RunT(t)
|
||||
st := store.NewWithClients(redis.NewClient(&redis.Options{Addr: m.Addr()}), nil)
|
||||
t.Cleanup(func() { st.Close() })
|
||||
return &Counter{Planner: newPlanner(false), Store: st}, m
|
||||
}
|
||||
|
||||
// Executes the plan: counts come back from Redis and the crossing increment
|
||||
// is flagged exactly once.
|
||||
func TestCounter_Out(t *testing.T) {
|
||||
c, m := newCounter(t)
|
||||
ctx := context.Background()
|
||||
// noreply@example.com has A/o=5: the fifth message reaches it.
|
||||
var last *Outcome
|
||||
for i := 1; i <= 6; i++ {
|
||||
out, err := c.CountOut(ctx, "tr", "noreply@example.com", "", "ext@gmail.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(out.Results) != 3 || out.Results[0].Count != int64(i) {
|
||||
t.Fatalf("message %d: results = %+v", i, out.Results)
|
||||
}
|
||||
if reached := out.Results[0].Reached; reached != (i == 5) {
|
||||
t.Errorf("message %d: A reached = %v", i, reached)
|
||||
}
|
||||
last = out
|
||||
}
|
||||
if v, _ := m.Get(store.Key("A", "noreply@example.com", day, "o")); v != "6" {
|
||||
t.Errorf("redis A/o = %q, want 6", v)
|
||||
}
|
||||
if v, _ := m.Get(store.Key("O", "7", day, "o")); v != "6" {
|
||||
t.Errorf("redis O/o = %q, want 6", v)
|
||||
}
|
||||
if last.Skipped != "" || last.Route != "out" {
|
||||
t.Errorf("outcome = %+v", last)
|
||||
}
|
||||
}
|
||||
|
||||
// After the web app resets a counter (DEL), the next crossing fires again.
|
||||
func TestCounter_ReachedAgainAfterReset(t *testing.T) {
|
||||
c, m := newCounter(t)
|
||||
ctx := context.Background()
|
||||
for i := 0; i < 5; i++ {
|
||||
if _, err := c.CountOut(ctx, "tr", "noreply@example.com", "", "ext@gmail.com"); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
}
|
||||
m.Del(store.Key("A", "noreply@example.com", day, "o"))
|
||||
var reachedAt int
|
||||
for i := 1; i <= 5; i++ {
|
||||
out, err := c.CountOut(ctx, "tr", "noreply@example.com", "", "ext@gmail.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if out.Results[0].Reached {
|
||||
reachedAt = i
|
||||
}
|
||||
}
|
||||
if reachedAt != 5 {
|
||||
t.Errorf("after reset the limit was reached at message %d, want 5", reachedAt)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCounter_In(t *testing.T) {
|
||||
c, m := newCounter(t)
|
||||
out, err := c.CountIn(context.Background(), "tr", "s@gmail.com", "y@example.com", false)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if len(out.Results) != 3 || out.Results[2].Count != 1 || out.Results[2].Level != "O" {
|
||||
t.Errorf("results = %+v", out.Results)
|
||||
}
|
||||
if v, _ := m.Get(store.Key("D", "example.com", day, "i")); v != "1" {
|
||||
t.Errorf("redis D/i = %q, want 1", v)
|
||||
}
|
||||
}
|
||||
|
||||
// A skipped plan must not touch Redis at all.
|
||||
func TestCounter_SkippedDoesNotTouchRedis(t *testing.T) {
|
||||
c, m := newCounter(t)
|
||||
out, err := c.CountIn(context.Background(), "tr", "support@migadu.com", "y@example.com", false)
|
||||
if err != nil || out.Skipped != SkipExempt || len(out.Results) != 0 {
|
||||
t.Fatalf("outcome = %+v, %v", out, err)
|
||||
}
|
||||
if keys := m.Keys(); len(keys) != 0 {
|
||||
t.Errorf("redis keys = %v, want none", keys)
|
||||
}
|
||||
}
|
||||
|
||||
func TestCounter_RedisDown(t *testing.T) {
|
||||
c, m := newCounter(t)
|
||||
m.Close()
|
||||
if _, err := c.CountIn(context.Background(), "tr", "s@gmail.com", "y@example.com", false); err == nil {
|
||||
t.Error("expected error when redis is down")
|
||||
}
|
||||
}
|
||||
|
||||
// --- Jobs with unusable sides (review #2, #3) --------------------------------
|
||||
|
||||
func TestPlan_OutSenderUnusableCountsRecipient(t *testing.T) {
|
||||
job := Job{Route: RouteOut, TraceID: "u1", Recipient: "y@example.com",
|
||||
SenderUnusable: Unusable{Role: "sender", Value: "nobody", Reason: "invalid e-mail address"}}
|
||||
plan, err := newPlanner(false).Plan(job)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "y@example.com", "i", "y@example.com", 500, true},
|
||||
{"D", "example.com", "i", "y@example.com", 10000, true},
|
||||
{"O", "7", "i", "y@example.com", 10000, true},
|
||||
})
|
||||
if len(plan.Unusable) != 1 || plan.Unusable[0].Role != "sender" || plan.Unusable[0].Value != "nobody" {
|
||||
t.Errorf("unusable = %+v", plan.Unusable)
|
||||
}
|
||||
if len(plan.Lookups) != 1 || plan.Lookups[0].Role != "recipient" {
|
||||
t.Errorf("only the recipient should be looked up: %+v", plan.Lookups)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPlan_OutRecipientUnusableCountsSender(t *testing.T) {
|
||||
job := Job{Route: RouteOut, TraceID: "u2", AuthUser: "x@example.com", LoginSuffix: "SUPPORTTOKEN", MailFrom: "x@example.com",
|
||||
RecipientUnusable: Unusable{Role: "recipient", Value: "u@[192.0.2.1]", Reason: "invalid e-mail address"}}
|
||||
plan, err := newPlanner(false).Plan(job)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "x@example.com", "o", "x@example.com", 100, true},
|
||||
{"D", "example.com", "o", "x@example.com", 2200, true},
|
||||
{"O", "7", "o", "x@example.com", 2200, true},
|
||||
})
|
||||
if plan.LoginSuffix != "SUPPORTTOKEN" || plan.AuthUser != "x@example.com" {
|
||||
t.Errorf("login suffix not carried: %+v", plan)
|
||||
}
|
||||
if len(plan.Unusable) != 1 || plan.Unusable[0].Role != "recipient" {
|
||||
t.Errorf("unusable = %+v", plan.Unusable)
|
||||
}
|
||||
}
|
||||
|
||||
// An unusable mail_from on /in is logged, treated as non-exempt, and the
|
||||
// recipient is counted.
|
||||
func TestPlan_InMailFromUnusable(t *testing.T) {
|
||||
job := Job{Route: RouteIn, TraceID: "u3", MailFrom: "john smith@x.com", Recipient: "y@example.com"}
|
||||
plan, err := newPlanner(false).Plan(job)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if plan.Skipped != "" || len(plan.Increments) != 3 || plan.MailFrom != "john smith@x.com" {
|
||||
t.Errorf("plan = %+v", plan)
|
||||
}
|
||||
}
|
||||
|
||||
// The exempt check still applies to usable sides of a Job.
|
||||
func TestPlan_ExemptViaJob(t *testing.T) {
|
||||
plan, err := newPlanner(false).Plan(Job{Route: RouteIn, MailFrom: "support@migadu.com", MailFromUsable: true, Recipient: "y@example.com"})
|
||||
if err != nil || plan.Skipped != SkipExempt {
|
||||
t.Errorf("plan = %+v, %v", plan, err)
|
||||
}
|
||||
}
|
||||
|
||||
// --- Execute: partial EXEC failure and lost replies (review #4) --------------
|
||||
|
||||
// One key holding garbage fails its own INCR; the others are counted and the
|
||||
// crossing among them is still detected.
|
||||
func TestCounter_PartialExecFailure(t *testing.T) {
|
||||
c, m := newCounter(t)
|
||||
m.Set(store.Key("O", "7", day, "o"), "garbage")
|
||||
m.Set(store.Key("A", "noreply@example.com", day, "o"), "4") // A/o=5: this increment crosses
|
||||
out, err := c.CountOut(context.Background(), "p1", "noreply@example.com", "", "ext@gmail.com")
|
||||
if err != nil {
|
||||
t.Fatalf("a per-key failure must not fail the request: %v", err)
|
||||
}
|
||||
if len(out.Results) != 3 {
|
||||
t.Fatalf("results = %+v", out.Results)
|
||||
}
|
||||
a, d, o := out.Results[0], out.Results[1], out.Results[2]
|
||||
if a.Err != nil || a.Count != 5 || !a.Reached {
|
||||
t.Errorf("A = %+v, want count 5 reached", a)
|
||||
}
|
||||
if d.Err != nil || d.Count != 1 {
|
||||
t.Errorf("D = %+v", d)
|
||||
}
|
||||
if o.Err == nil || o.Reached {
|
||||
t.Errorf("O = %+v, want its own error and no crossing", o)
|
||||
}
|
||||
if out.Unknown != nil {
|
||||
t.Errorf("outcome is known: %v", out.Unknown)
|
||||
}
|
||||
}
|
||||
|
||||
// delayProxy forwards TCP to backend and, once it has seen EXEC in the
|
||||
// request stream, holds back the next `delayed` replies by `delay`, so the
|
||||
// client's read timeout fires after Redis has applied the transaction (the
|
||||
// connection handshake replies pass undelayed).
|
||||
func delayProxy(t *testing.T, backend string, delayed int, delay time.Duration) string {
|
||||
ln, err := net.Listen("tcp", "127.0.0.1:0")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
t.Cleanup(func() { ln.Close() })
|
||||
var remaining atomic.Int64
|
||||
remaining.Store(int64(delayed))
|
||||
go func() {
|
||||
for {
|
||||
c, err := ln.Accept()
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
b, err := net.Dial("tcp", backend)
|
||||
if err != nil {
|
||||
c.Close()
|
||||
continue
|
||||
}
|
||||
var armed atomic.Bool
|
||||
go func() {
|
||||
defer b.Close()
|
||||
buf := make([]byte, 32*1024)
|
||||
for {
|
||||
n, err := c.Read(buf)
|
||||
if n > 0 {
|
||||
if bytes.Contains(bytes.ToLower(buf[:n]), []byte("exec")) {
|
||||
armed.Store(true)
|
||||
}
|
||||
if _, werr := b.Write(buf[:n]); werr != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer c.Close()
|
||||
buf := make([]byte, 32*1024)
|
||||
for {
|
||||
n, err := b.Read(buf)
|
||||
if n > 0 {
|
||||
if armed.Load() && remaining.Add(-1) >= 0 {
|
||||
time.Sleep(delay)
|
||||
}
|
||||
if _, werr := c.Write(buf[:n]); werr != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
}()
|
||||
return ln.Addr().String()
|
||||
}
|
||||
|
||||
// The EXEC reply is lost after Redis applied the increments: the counter is
|
||||
// re-read from the master and the crossing is still reported.
|
||||
func TestCounter_LostReplyReReadsMaster(t *testing.T) {
|
||||
m := miniredis.RunT(t)
|
||||
proxy := delayProxy(t, m.Addr(), 1, 800*time.Millisecond)
|
||||
st := store.NewWithClients(redis.NewClient(&redis.Options{Addr: proxy, ReadTimeout: 200 * time.Millisecond, PoolSize: 2}), nil)
|
||||
t.Cleanup(func() { st.Close() })
|
||||
c := &Counter{Planner: newPlanner(false), Store: st}
|
||||
m.Set(store.Key("A", "noreply@example.com", day, "o"), "4") // A/o=5
|
||||
|
||||
out, err := c.CountOut(context.Background(), "lost", "noreply@example.com", "", "ext@gmail.com")
|
||||
if err != nil {
|
||||
t.Fatalf("lost reply must be resolved by a re-read, got error: %v", err)
|
||||
}
|
||||
if out.Unknown == nil {
|
||||
t.Fatal("outcome should be flagged unknown (reply lost)")
|
||||
}
|
||||
if v, _ := m.Get(store.Key("A", "noreply@example.com", day, "o")); v != "5" {
|
||||
t.Fatalf("redis A/o = %q, the increment should have been applied", v)
|
||||
}
|
||||
if a := out.Results[0]; a.Count != 5 || !a.Reached {
|
||||
t.Errorf("re-read A = %+v, want count 5 reached", a)
|
||||
}
|
||||
if o := out.Results[2]; o.Count != 1 || o.Reached {
|
||||
t.Errorf("re-read O = %+v", o)
|
||||
}
|
||||
}
|
||||
|
||||
// A definite failure (connection refused) is still an error: nothing counted.
|
||||
func TestCounter_RefusedIsNotReRead(t *testing.T) {
|
||||
c, m := newCounter(t)
|
||||
m.Close()
|
||||
out, err := c.CountIn(context.Background(), "tr", "s@gmail.com", "y@example.com", false)
|
||||
if err == nil {
|
||||
t.Errorf("expected error, got outcome %+v", out)
|
||||
}
|
||||
}
|
||||
|
||||
// --- review #7: plus-addressing counts against the mailbox -------------------
|
||||
|
||||
func TestPlan_PlusTagCountsMailbox(t *testing.T) {
|
||||
p := newPlanner(false)
|
||||
plan, err := p.PlanOut("pt", "noreply+news@example.com", "", "y+1@example.com")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "noreply@example.com", "o", "noreply@example.com", 5, true}, // override noreply@/o still applies
|
||||
{"D", "example.com", "o", "noreply@example.com", 2200, true},
|
||||
{"O", "7", "o", "noreply@example.com", 2200, true},
|
||||
{"A", "y@example.com", "i", "y@example.com", 500, true},
|
||||
{"D", "example.com", "i", "y@example.com", 10000, true},
|
||||
{"O", "7", "i", "y@example.com", 10000, true},
|
||||
})
|
||||
if plan.Recipient != "y+1@example.com" || plan.AuthUser != "noreply+news@example.com" {
|
||||
t.Errorf("the request line keeps the addresses as sent: %+v", plan)
|
||||
}
|
||||
// Exempt applies to the mailbox too.
|
||||
plan, _ = p.PlanIn("pt", "s@gmail.com", "support+ticket@migadu.com", false)
|
||||
if plan.Skipped != SkipExempt {
|
||||
t.Errorf("support+ticket@ should be exempt: %+v", plan)
|
||||
}
|
||||
}
|
||||
|
||||
// --- review #8: a malformed CDB row is "not hosted", not a failure -----------
|
||||
|
||||
func TestPlan_MalformedRowIsNotHosted(t *testing.T) {
|
||||
db, err := domains.Open("../../review/_proof_tests/b/bad.cdb")
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
defer db.Close()
|
||||
p := &Planner{Lookup: db, Now: func() time.Time { return fixedNow }}
|
||||
plan, err := p.PlanOut("m", "a@good.com", "", "b@bad.com")
|
||||
if err != nil {
|
||||
t.Fatalf("malformed recipient row must not fail the request: %v", err)
|
||||
}
|
||||
checkIncrements(t, plan.Increments, []want{
|
||||
{"A", "a@good.com", "o", "a@good.com", 5, true},
|
||||
{"D", "good.com", "o", "a@good.com", 0, false},
|
||||
{"O", "7", "o", "a@good.com", 0, false},
|
||||
})
|
||||
if len(plan.Lookups) != 2 || plan.Lookups[1].Domain != nil || !errors.Is(plan.Lookups[1].Err, domains.ErrMalformed) {
|
||||
t.Errorf("recipient lookup should record the malformed row: %+v", plan.Lookups)
|
||||
}
|
||||
plan, err = p.PlanIn("m", "s@gmail.com", "b@bad.com", false)
|
||||
if err != nil || plan.Skipped != SkipUnknownDomain || len(plan.Increments) != 0 {
|
||||
t.Errorf("/in to a malformed row: plan=%+v err=%v, want skipped unknown_domain", plan, err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,138 @@
|
||||
package counting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"runtime/debug"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
)
|
||||
|
||||
// QueueConfig sizes the counting queue.
|
||||
type QueueConfig struct {
|
||||
Size int // queued jobs before Enqueue refuses (default 10 000)
|
||||
Workers int // concurrent workers (default 8)
|
||||
Timeout time.Duration // budget per job: plan, Redis, log (default 4s)
|
||||
}
|
||||
|
||||
// Observer receives queue events for metrics and logging; nil funcs are ignored.
|
||||
type Observer struct {
|
||||
Depth func(n int) // queue depth after every enqueue/dequeue
|
||||
Wait func(d time.Duration) // time a job waited in the queue
|
||||
Panic func(job Job, p any, stack []byte) // a handler panicked (recovered)
|
||||
}
|
||||
|
||||
// Queue decouples counting from the HTTP request: /in and /out validate,
|
||||
// enqueue a Job and answer 202; workers run the handler with their own
|
||||
// bounded context, so a client that hangs up or times out can no longer
|
||||
// cancel a count (review #1).
|
||||
type Queue struct {
|
||||
cfg QueueConfig
|
||||
jobs chan Job
|
||||
handler func(ctx context.Context, job Job)
|
||||
obs Observer
|
||||
|
||||
pending atomic.Int64 // enqueued and not yet finished
|
||||
mu sync.RWMutex // Enqueue vs Drain (sending on a closed channel panics)
|
||||
closed bool
|
||||
wg sync.WaitGroup
|
||||
}
|
||||
|
||||
// NewQueue applies defaults and starts the workers.
|
||||
func NewQueue(cfg QueueConfig, handler func(ctx context.Context, job Job), obs Observer) *Queue {
|
||||
if cfg.Size <= 0 {
|
||||
cfg.Size = 10000
|
||||
}
|
||||
if cfg.Workers <= 0 {
|
||||
cfg.Workers = 8
|
||||
}
|
||||
if cfg.Timeout <= 0 {
|
||||
cfg.Timeout = 4 * time.Second
|
||||
}
|
||||
q := &Queue{cfg: cfg, jobs: make(chan Job, cfg.Size), handler: handler, obs: obs}
|
||||
for i := 0; i < cfg.Workers; i++ {
|
||||
q.wg.Add(1)
|
||||
go q.worker()
|
||||
}
|
||||
return q
|
||||
}
|
||||
|
||||
// Enqueue hands a job to the workers without blocking. It returns false when
|
||||
// the queue is full or draining; the caller answers 503 and logs the drop.
|
||||
func (q *Queue) Enqueue(job Job) bool {
|
||||
if job.Received.IsZero() {
|
||||
job.Received = time.Now()
|
||||
}
|
||||
q.mu.RLock()
|
||||
defer q.mu.RUnlock()
|
||||
if q.closed {
|
||||
return false
|
||||
}
|
||||
select {
|
||||
case q.jobs <- job:
|
||||
q.pending.Add(1)
|
||||
q.observeDepth()
|
||||
return true
|
||||
default:
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Depth is the number of jobs waiting for a worker.
|
||||
func (q *Queue) Depth() int { return len(q.jobs) }
|
||||
|
||||
// Pending is the number of jobs enqueued and not yet finished (waiting or
|
||||
// in flight); tests and the shutdown log use it.
|
||||
func (q *Queue) Pending() int { return int(q.pending.Load()) }
|
||||
|
||||
// Drain stops accepting, lets the workers finish every queued job and waits
|
||||
// for them, bounded by ctx. On an expired ctx it reports how many jobs were
|
||||
// still pending.
|
||||
func (q *Queue) Drain(ctx context.Context) error {
|
||||
q.mu.Lock()
|
||||
if !q.closed {
|
||||
q.closed = true
|
||||
close(q.jobs)
|
||||
}
|
||||
q.mu.Unlock()
|
||||
done := make(chan struct{})
|
||||
go func() { q.wg.Wait(); close(done) }()
|
||||
select {
|
||||
case <-done:
|
||||
return nil
|
||||
case <-ctx.Done():
|
||||
return fmt.Errorf("counting queue: %d jobs still pending: %w", q.Pending(), ctx.Err())
|
||||
}
|
||||
}
|
||||
|
||||
func (q *Queue) worker() {
|
||||
defer q.wg.Done()
|
||||
for job := range q.jobs {
|
||||
if q.obs.Wait != nil {
|
||||
q.obs.Wait(time.Since(job.Received))
|
||||
}
|
||||
q.run(job)
|
||||
q.pending.Add(-1)
|
||||
q.observeDepth()
|
||||
}
|
||||
}
|
||||
|
||||
// run executes one job under the queue's own timeout, never a request
|
||||
// context, and keeps the worker alive when the handler panics.
|
||||
func (q *Queue) run(job Job) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), q.cfg.Timeout)
|
||||
defer cancel()
|
||||
defer func() {
|
||||
if p := recover(); p != nil && q.obs.Panic != nil {
|
||||
q.obs.Panic(job, p, debug.Stack())
|
||||
}
|
||||
}()
|
||||
q.handler(ctx, job)
|
||||
}
|
||||
|
||||
func (q *Queue) observeDepth() {
|
||||
if q.obs.Depth != nil {
|
||||
q.obs.Depth(len(q.jobs))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,179 @@
|
||||
package counting
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
type recorder struct {
|
||||
mu sync.Mutex
|
||||
jobs []Job
|
||||
waits []time.Duration
|
||||
depths []int
|
||||
panics int
|
||||
}
|
||||
|
||||
func (r *recorder) handle(ctx context.Context, job Job) {
|
||||
r.mu.Lock()
|
||||
r.jobs = append(r.jobs, job)
|
||||
r.mu.Unlock()
|
||||
}
|
||||
|
||||
func (r *recorder) observer() Observer {
|
||||
return Observer{
|
||||
Depth: func(n int) { r.mu.Lock(); r.depths = append(r.depths, n); r.mu.Unlock() },
|
||||
Wait: func(d time.Duration) { r.mu.Lock(); r.waits = append(r.waits, d); r.mu.Unlock() },
|
||||
Panic: func(job Job, p any, stack []byte) { r.mu.Lock(); r.panics++; r.mu.Unlock() },
|
||||
}
|
||||
}
|
||||
|
||||
func (r *recorder) count() int {
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
return len(r.jobs)
|
||||
}
|
||||
|
||||
func settle(t *testing.T, q *Queue) {
|
||||
t.Helper()
|
||||
deadline := time.Now().Add(5 * time.Second)
|
||||
for q.Pending() > 0 {
|
||||
if time.Now().After(deadline) {
|
||||
t.Fatalf("queue did not settle: pending=%d", q.Pending())
|
||||
}
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueue_ProcessesJobs(t *testing.T) {
|
||||
r := &recorder{}
|
||||
q := NewQueue(QueueConfig{Size: 10, Workers: 2, Timeout: time.Second}, r.handle, r.observer())
|
||||
defer q.Drain(context.Background())
|
||||
for i := 0; i < 5; i++ {
|
||||
if !q.Enqueue(Job{Route: RouteIn, TraceID: "t", Recipient: "y@example.com", Received: time.Now()}) {
|
||||
t.Fatalf("enqueue %d refused", i)
|
||||
}
|
||||
}
|
||||
settle(t, q)
|
||||
if r.count() != 5 {
|
||||
t.Errorf("handled %d jobs, want 5", r.count())
|
||||
}
|
||||
r.mu.Lock()
|
||||
defer r.mu.Unlock()
|
||||
if len(r.waits) != 5 || len(r.depths) == 0 {
|
||||
t.Errorf("observer: waits=%d depths=%d", len(r.waits), len(r.depths))
|
||||
}
|
||||
for _, w := range r.waits {
|
||||
if w < 0 || w > time.Second {
|
||||
t.Errorf("queue wait %v out of range", w)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Workers run the handler with their own bounded context, never a request's.
|
||||
func TestQueue_HandlerContextHasTimeout(t *testing.T) {
|
||||
var got atomic.Value
|
||||
q := NewQueue(QueueConfig{Size: 1, Workers: 1, Timeout: 250 * time.Millisecond}, func(ctx context.Context, job Job) {
|
||||
dl, ok := ctx.Deadline()
|
||||
got.Store([2]any{ok, time.Until(dl)})
|
||||
}, Observer{})
|
||||
defer q.Drain(context.Background())
|
||||
q.Enqueue(Job{})
|
||||
settle(t, q)
|
||||
v := got.Load().([2]any)
|
||||
if !v[0].(bool) || v[1].(time.Duration) > 250*time.Millisecond || v[1].(time.Duration) < 100*time.Millisecond {
|
||||
t.Errorf("deadline = %v", v)
|
||||
}
|
||||
}
|
||||
|
||||
// A full queue refuses the job instead of blocking the HTTP handler.
|
||||
func TestQueue_FullRefuses(t *testing.T) {
|
||||
release := make(chan struct{})
|
||||
q := NewQueue(QueueConfig{Size: 2, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) { <-release }, Observer{})
|
||||
// One job occupies the worker ...
|
||||
if !q.Enqueue(Job{}) {
|
||||
t.Fatal("first enqueue refused")
|
||||
}
|
||||
deadline := time.Now().Add(time.Second)
|
||||
for q.Depth() != 0 && time.Now().Before(deadline) {
|
||||
time.Sleep(time.Millisecond)
|
||||
}
|
||||
// ... two fill the channel.
|
||||
for i := 0; i < 2; i++ {
|
||||
if !q.Enqueue(Job{}) {
|
||||
t.Fatalf("enqueue %d should be accepted", i)
|
||||
}
|
||||
}
|
||||
if q.Enqueue(Job{}) {
|
||||
t.Error("enqueue on a full queue should be refused")
|
||||
}
|
||||
if q.Depth() != 2 || q.Pending() != 3 {
|
||||
t.Errorf("depth=%d pending=%d, want 2/3", q.Depth(), q.Pending())
|
||||
}
|
||||
close(release)
|
||||
if err := q.Drain(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if q.Enqueue(Job{}) {
|
||||
t.Error("enqueue after Drain must be refused, not panic")
|
||||
}
|
||||
}
|
||||
|
||||
// Drain waits for queued and in-flight jobs; a context that expires first
|
||||
// reports what was left behind.
|
||||
func TestQueue_Drain(t *testing.T) {
|
||||
r := &recorder{}
|
||||
slow := func(ctx context.Context, job Job) { time.Sleep(50 * time.Millisecond); r.handle(ctx, job) }
|
||||
q := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, slow, Observer{})
|
||||
for i := 0; i < 4; i++ {
|
||||
q.Enqueue(Job{})
|
||||
}
|
||||
start := time.Now()
|
||||
if err := q.Drain(context.Background()); err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
if r.count() != 4 || time.Since(start) < 150*time.Millisecond {
|
||||
t.Errorf("Drain returned after %v with %d jobs handled, want all 4", time.Since(start), r.count())
|
||||
}
|
||||
|
||||
block := make(chan struct{})
|
||||
q2 := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) { <-block }, Observer{})
|
||||
q2.Enqueue(Job{})
|
||||
q2.Enqueue(Job{})
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
err := q2.Drain(ctx)
|
||||
if err == nil || !errors.Is(err, context.DeadlineExceeded) {
|
||||
t.Errorf("Drain with a stuck worker should report the deadline: %v", err)
|
||||
}
|
||||
close(block)
|
||||
}
|
||||
|
||||
// A panicking handler is recovered and reported; the worker keeps serving.
|
||||
func TestQueue_RecoversPanic(t *testing.T) {
|
||||
r := &recorder{}
|
||||
q := NewQueue(QueueConfig{Size: 10, Workers: 1, Timeout: time.Second}, func(ctx context.Context, job Job) {
|
||||
if job.TraceID == "boom" {
|
||||
panic("boom")
|
||||
}
|
||||
r.handle(ctx, job)
|
||||
}, r.observer())
|
||||
defer q.Drain(context.Background())
|
||||
q.Enqueue(Job{TraceID: "boom"})
|
||||
q.Enqueue(Job{TraceID: "ok"})
|
||||
settle(t, q)
|
||||
if r.count() != 1 || r.panics != 1 {
|
||||
t.Errorf("handled=%d panics=%d, want 1/1", r.count(), r.panics)
|
||||
}
|
||||
}
|
||||
|
||||
func TestQueue_Defaults(t *testing.T) {
|
||||
q := NewQueue(QueueConfig{}, func(context.Context, Job) {}, Observer{})
|
||||
defer q.Drain(context.Background())
|
||||
if q.cfg.Size != 10000 || q.cfg.Workers != 8 || q.cfg.Timeout != 4*time.Second {
|
||||
t.Errorf("defaults = %+v", q.cfg)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user