Improve config reload consistency, add tests

Author: foxcpp

framework/dns/debugflags.go

@@ -32,5 +32,12 @@ func init() {
 		Usage:       "replace the DNS resolver address",
 		Value:       "system-default",
 		Destination: &overrideServ,
+		Action: func(context *cli.Context, s string) error {
+			if s != "" && s != "system-default" {
+				override(s)
+			}
+			overrideServ = s
+			return nil
+		},
 	})
 }

framework/dns/resolver.go

@@ -53,9 +53,5 @@ func LookupAddr(ctx context.Context, r Resolver, ip net.IP) (string, error) {
 }
 
 func DefaultResolver() Resolver {
-	if overrideServ != "" && overrideServ != "system-default" {
-		override(overrideServ)
-	}
-
 	return net.DefaultResolver
 }

framework/module/lifetime.go

@@ -37,18 +37,27 @@ type ReloadModule interface {
 	Reload() error
 }
 
+// EarlyStopModule is a LifetimeModule that needs to do some bookkeeping
+// before new server instance starts during reload.
+type EarlyStopModule interface {
+	LifetimeModule
+	EarlyStop() error
+}
+
 type LifetimeTracker struct {
 	logger    *log.Logger
 	instances []*struct {
-		mod     LifetimeModule
-		started bool
+		mod          LifetimeModule
+		started      bool
+		earlyStopped bool
 	}
 }
 
 func (lt *LifetimeTracker) Add(mod LifetimeModule) {
 	lt.instances = append(lt.instances, &struct {
-		mod     LifetimeModule
-		started bool
+		mod          LifetimeModule
+		started      bool
+		earlyStopped bool
 	}{mod: mod, started: false})
 }
 
@@ -59,6 +68,9 @@ func (lt *LifetimeTracker) StartAll() error {
 			continue
 		}
 
+		lt.logger.DebugMsg("starting module",
+			"mod_name", entry.mod.Name(), "inst_name", entry.mod.InstanceName())
+
 		if err := entry.mod.Start(); err != nil {
 			if err := lt.StopAll(); err != nil {
 				lt.logger.Error("StopAll failed after Start fail", err)
@@ -96,6 +108,32 @@ func (lt *LifetimeTracker) ReloadAll() error {
 	return nil
 }
 
+func (lt *LifetimeTracker) EarlyStopAll() error {
+	for i := len(lt.instances) - 1; i >= 0; i-- {
+		entry := lt.instances[i]
+
+		if !entry.started {
+			continue
+		}
+
+		rsm, ok := entry.mod.(EarlyStopModule)
+		if !ok {
+			continue
+		}
+
+		if err := rsm.EarlyStop(); err != nil {
+			lt.logger.Error("module early stop failed", err,
+				"mod_name", entry.mod.Name(), "inst_name", entry.mod.InstanceName())
+			continue
+		}
+		lt.logger.DebugMsg("module early stopped",
+			"mod_name", entry.mod.Name(), "inst_name", entry.mod.InstanceName())
+
+		entry.earlyStopped = true
+	}
+	return nil
+}
+
 // StopAll calls Stop for all registered LifetimeModule instances.
 func (lt *LifetimeTracker) StopAll() error {
 	for i := len(lt.instances) - 1; i >= 0; i-- {

internal/sqlite/modernc_sqlite3.go

@@ -1,4 +1,4 @@
-//go:build !nosqlite3 && !cgo
+//go:build (!nosqlite3 && !cgo) || modernc
 
 /*
 Maddy Mail Server - Composable all-in-one email server.

internal/sqlite/sqlite3.go

@@ -1,4 +1,4 @@
-//go:build !nosqlite3 && cgo
+//go:build !nosqlite3 && cgo && !modernc
 
 /*
 Maddy Mail Server - Composable all-in-one email server.

internal/target/queue/queue.go

@@ -122,7 +122,7 @@ type Queue struct {
 	location         string
 	hostname         string
 	autogenMsgDomain string
-	wheel            *TimeWheel
+	wheel            *TimeWheel[queueSlot]
 
 	dsnPipeline module.DeliveryTarget
 
@@ -211,6 +211,9 @@ func (q *Queue) Configure(inlineArgs []string, cfg *config.Map) error {
 	cfg.Bool("debug", true, false, &q.Log.Debug)
 	cfg.Int("max_tries", false, false, 20, &q.maxTries)
 	cfg.Int("max_parallelism", false, false, 16, &q.maxParallelism)
+	cfg.Duration("post_init_delay", false, false, q.postInitDelay, &q.postInitDelay)
+	cfg.Duration("initial_retry_time", false, false, q.initialRetryTime, &q.initialRetryTime)
+	cfg.Float("retry_time_scale", false, false, q.retryTimeScale, &q.retryTimeScale)
 	cfg.String("location", false, false, q.location, &q.location)
 	cfg.Custom("target", false, true, nil, modconfig.DeliveryDirective, &q.Target)
 	cfg.String("hostname", true, true, "", &q.hostname)
@@ -249,7 +252,7 @@ func (q *Queue) Start() error {
 }
 
 func (q *Queue) start(maxParallelism int) error {
-	q.wheel = NewTimeWheel(q.dispatch)
+	q.wheel = NewTimeWheel[queueSlot](q.dispatch)
 	q.deliverySemaphore = make(chan struct{}, maxParallelism)
 
 	if err := q.readDiskQueue(); err != nil {
@@ -261,13 +264,18 @@ func (q *Queue) start(maxParallelism int) error {
 	return nil
 }
 
-func (q *Queue) Stop() error {
+func (q *Queue) EarlyStop() error {
+	// We must ensure queue state is consistent on disk before we proceed
+	// with configuration reload.
 	q.wheel.Close()
 	q.deliveryWg.Wait()
-
 	return nil
 }
 
+func (q *Queue) Stop() error {
+	return q.EarlyStop()
+}
+
 // discardBroken changes the name of metadata file to have .meta_broken
 // extension.
 //
@@ -283,8 +291,8 @@ func (q *Queue) discardBroken(id string) {
 	}
 }
 
-func (q *Queue) dispatch(value TimeSlot) {
-	slot := value.Value.(queueSlot)
+func (q *Queue) dispatch(ctx context.Context, value TimeSlot[queueSlot]) {
+	slot := value.Value
 
 	q.Log.Debugln("starting delivery for", slot.ID)
 
@@ -329,7 +337,7 @@ func (q *Queue) dispatch(value TimeSlot) {
 			body = slot.Body
 		}
 
-		q.tryDelivery(meta, hdr, body)
+		q.tryDelivery(ctx, meta, hdr, body)
 	}()
 }
 
@@ -373,10 +381,10 @@ func toSMTPErr(err error) *smtp.SMTPError {
 	return res
 }
 
-func (q *Queue) tryDelivery(meta *QueueMetadata, header textproto.Header, body buffer.Buffer) {
+func (q *Queue) tryDelivery(ctx context.Context, meta *QueueMetadata, header textproto.Header, body buffer.Buffer) {
 	dl := target.DeliveryLogger(q.Log, meta.MsgMeta)
 
-	partialErr := q.deliver(meta, header, body)
+	partialErr := q.deliver(ctx, meta, header, body)
 	dl.Debugf("errors: %v", partialErr.Errs)
 
 	// While iterating the list of recipients we also pick the smallest tries count
@@ -460,7 +468,7 @@ func (q *Queue) tryDelivery(meta *QueueMetadata, header textproto.Header, body b
 	})
 }
 
-func (q *Queue) deliver(meta *QueueMetadata, header textproto.Header, body buffer.Buffer) partialError {
+func (q *Queue) deliver(ctx context.Context, meta *QueueMetadata, header textproto.Header, body buffer.Buffer) partialError {
 	dl := target.DeliveryLogger(q.Log, meta.MsgMeta)
 	perr := partialError{
 		Errs:       map[string]error{},
@@ -471,7 +479,7 @@ func (q *Queue) deliver(meta *QueueMetadata, header textproto.Header, body buffe
 	msgMeta.ID = msgMeta.ID + "-" + strconv.FormatInt(time.Now().Unix(), 16)
 	dl.Debugf("using message ID = %s", msgMeta.ID)
 
-	msgCtx, msgTask := trace.NewTask(context.Background(), "Queue delivery")
+	msgCtx, msgTask := trace.NewTask(ctx, "Queue delivery")
 	defer msgTask.End()
 
 	mailCtx, mailTask := trace.NewTask(msgCtx, "MAIL FROM")
@@ -486,6 +494,15 @@ func (q *Queue) deliver(meta *QueueMetadata, header textproto.Header, body buffe
 	}
 	dl.Debugf("target.StartDelivery OK")
 
+	// Check in case delivery implementation is actually
+	// context-unaware.
+	if err := mailCtx.Err(); err != nil {
+		for _, rcpt := range meta.To {
+			perr.Errs[rcpt] = err
+		}
+		return perr
+	}
+
 	var acceptedRcpts []string
 	for _, rcpt := range meta.To {
 		rcptCtx, rcptTask := trace.NewTask(msgCtx, "RCPT TO")
@@ -497,6 +514,15 @@ func (q *Queue) deliver(meta *QueueMetadata, header textproto.Header, body buffe
 			acceptedRcpts = append(acceptedRcpts, rcpt)
 		}
 		rcptTask.End()
+
+		// Check in case delivery implementation is actually
+		// context-unaware.
+		if err := mailCtx.Err(); err != nil {
+			for _, rcpt := range meta.To {
+				perr.Errs[rcpt] = err
+			}
+			return perr
+		}
 	}
 
 	if len(acceptedRcpts) == 0 {
@@ -513,6 +539,10 @@ func (q *Queue) deliver(meta *QueueMetadata, header textproto.Header, body buffe
 		}
 	}
 
+	// At this point, it is too late to abort delivery. We should complete
+	// it or fail it consistently.
+	msgCtx = context.WithoutCancel(msgCtx)
+
 	bodyCtx, bodyTask := trace.NewTask(msgCtx, "DATA")
 	defer bodyTask.End()
 

internal/target/queue/timewheel.go

@@ -20,64 +20,69 @@ package queue
 
 import (
 	"container/list"
+	"context"
 	"sync"
 	"sync/atomic"
 	"time"
 )
 
-type TimeSlot struct {
+type TimeSlot[Value any] struct {
 	Time  time.Time
-	Value interface{}
+	Value Value
 }
 
-type TimeWheel struct {
+type TimeWheel[Value any] struct {
 	stopped uint32
 
 	slots     *list.List
 	slotsLock sync.Mutex
 
 	updateNotify chan time.Time
 	stopNotify   chan struct{}
+	tickerCtx    context.Context
+	tickerCancel context.CancelFunc
 
-	dispatch func(TimeSlot)
+	dispatch func(context.Context, TimeSlot[Value])
 }
 
-func NewTimeWheel(dispatch func(TimeSlot)) *TimeWheel {
-	tw := &TimeWheel{
+func NewTimeWheel[Value any](dispatch func(context.Context, TimeSlot[Value])) *TimeWheel[Value] {
+	ctx, cancel := context.WithCancel(context.Background())
+
+	tw := &TimeWheel[Value]{
 		slots:        list.New(),
 		stopNotify:   make(chan struct{}),
+		tickerCtx:    ctx,
+		tickerCancel: cancel,
 		updateNotify: make(chan time.Time),
 		dispatch:     dispatch,
 	}
-	go tw.tick()
+	go tw.tick(context.Background())
 	return tw
 }
 
-func (tw *TimeWheel) Add(target time.Time, value interface{}) {
+func (tw *TimeWheel[Value]) Add(target time.Time, value Value) {
 	if atomic.LoadUint32(&tw.stopped) == 1 {
 		// Already stopped, ignore.
 		return
 	}
 
-	if value == nil {
-		panic("can't insert nil objects into TimeWheel queue")
-	}
-
 	tw.slotsLock.Lock()
-	tw.slots.PushBack(TimeSlot{Time: target, Value: value})
+	tw.slots.PushBack(TimeSlot[Value]{Time: target, Value: value})
 	tw.slotsLock.Unlock()
 
 	tw.updateNotify <- target
 }
 
-func (tw *TimeWheel) Close() {
+func (tw *TimeWheel[Value]) Close() {
 	atomic.StoreUint32(&tw.stopped, 1)
 
 	// Idempotent Close is convenient sometimes.
 	if tw.stopNotify == nil {
 		return
 	}
 
+	tw.tickerCancel()
+
 	tw.stopNotify <- struct{}{}
 	<-tw.stopNotify
 
@@ -86,16 +91,16 @@ func (tw *TimeWheel) Close() {
 	close(tw.updateNotify)
 }
 
-func (tw *TimeWheel) tick() {
+func (tw *TimeWheel[Value]) tick(ctx context.Context) {
 	for {
 		now := time.Now()
 		// Look for list element closest to now.
 		tw.slotsLock.Lock()
-		var closestSlot TimeSlot
+		var closestSlot TimeSlot[Value]
 		var closestEl *list.Element
 		for e := tw.slots.Front(); e != nil; e = e.Next() {
-			slot := e.Value.(TimeSlot)
-			if slot.Time.Sub(now) < closestSlot.Time.Sub(now) || closestSlot.Value == nil {
+			slot := e.Value.(TimeSlot[Value])
+			if slot.Time.Sub(now) < closestSlot.Time.Sub(now) || closestEl == nil {
 				closestSlot = slot
 				closestEl = e
 			}
@@ -124,7 +129,7 @@ func (tw *TimeWheel) tick() {
 				tw.slots.Remove(closestEl)
 				tw.slotsLock.Unlock()
 
-				tw.dispatch(closestSlot)
+				tw.dispatch(ctx, closestSlot)
 
 				break selectloop
 			case newTarget := <-tw.updateNotify: