538 lines
9.8 KiB
Go
538 lines
9.8 KiB
Go
package prism
|
|
|
|
import (
|
|
"container/list"
|
|
"context"
|
|
"fmt"
|
|
"git.wisehodl.dev/jay/go-mana-component"
|
|
"log/slog"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Types
|
|
// ----------------------------------------------------------------------------
|
|
|
|
// Letters
|
|
|
|
type LetterID = uint64
|
|
|
|
type OutboundLetter struct {
|
|
id uint64
|
|
peerID string
|
|
data Envelope
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
}
|
|
|
|
type LetterOutcomeKind int
|
|
|
|
const (
|
|
OutcomeSent LetterOutcomeKind = iota
|
|
OutcomeExpired
|
|
OutcomeCancelled
|
|
OutcomeRejected
|
|
)
|
|
|
|
func (k LetterOutcomeKind) String() string {
|
|
switch k {
|
|
case OutcomeSent:
|
|
return "sent"
|
|
case OutcomeExpired:
|
|
return "expired"
|
|
case OutcomeCancelled:
|
|
return "cancelled"
|
|
case OutcomeRejected:
|
|
return "rejected"
|
|
default:
|
|
return "unknown"
|
|
}
|
|
}
|
|
|
|
type LetterOutcome struct {
|
|
LetterID uint64
|
|
PeerID string
|
|
Kind LetterOutcomeKind
|
|
SentAt time.Time
|
|
MissedAt time.Time
|
|
Retries int
|
|
}
|
|
|
|
// Postmaster
|
|
|
|
type Postmaster struct {
|
|
couriers map[string]*Courier
|
|
poolHasPeer func(id string) (string, bool)
|
|
poolEvents <-chan PoolEvent // Adapter.Subscribe
|
|
poolSend PoolSendFunc // Adapter.Send
|
|
counter atomic.Uint64
|
|
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
mu sync.RWMutex
|
|
wg sync.WaitGroup
|
|
cfg postmasterConfig
|
|
handler slog.Handler
|
|
logger *slog.Logger
|
|
}
|
|
|
|
// Courier
|
|
|
|
type Courier struct {
|
|
cmd chan courierCommand
|
|
sendFunc func(data Envelope) error
|
|
|
|
// state
|
|
queue list.List
|
|
connected bool
|
|
sending bool
|
|
|
|
ctx context.Context
|
|
cancel context.CancelFunc
|
|
mu sync.Mutex
|
|
wg sync.WaitGroup
|
|
logger *slog.Logger
|
|
}
|
|
|
|
// Commands
|
|
|
|
type courierCommand interface {
|
|
apply(c *Courier)
|
|
}
|
|
|
|
// Options
|
|
|
|
const (
|
|
DefaultPostmasterDeadline = 30 * time.Second
|
|
)
|
|
|
|
type PostmasterOption func(*postmasterConfig)
|
|
|
|
type postmasterConfig struct {
|
|
defaultDeadline time.Duration
|
|
}
|
|
|
|
func WithDefaultDeadline(d time.Duration) PostmasterOption {
|
|
return func(c *postmasterConfig) { c.defaultDeadline = d }
|
|
}
|
|
|
|
type SendOption func(*sendConfig)
|
|
|
|
type sendConfig struct {
|
|
deadline time.Duration
|
|
}
|
|
|
|
func WithDeadline(d time.Duration) SendOption {
|
|
return func(c *sendConfig) { c.deadline = d }
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Postmaster
|
|
// ----------------------------------------------------------------------------
|
|
|
|
func NewPostmaster(
|
|
ctx context.Context,
|
|
poolHasPeer func(id string) (string, bool),
|
|
poolEvents <-chan PoolEvent,
|
|
poolSendFunc PoolSendFunc,
|
|
handler slog.Handler,
|
|
opts ...PostmasterOption,
|
|
) *Postmaster {
|
|
ctx, cancel := context.WithCancel(
|
|
component.MustNew(ctx, "prism", "postmaster"))
|
|
|
|
cfg := postmasterConfig{
|
|
defaultDeadline: DefaultPostmasterDeadline,
|
|
}
|
|
for _, opt := range opts {
|
|
opt(&cfg)
|
|
}
|
|
|
|
pm := &Postmaster{
|
|
couriers: make(map[string]*Courier),
|
|
poolHasPeer: poolHasPeer,
|
|
poolEvents: poolEvents,
|
|
poolSend: poolSendFunc,
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
cfg: cfg,
|
|
}
|
|
|
|
if handler != nil {
|
|
comp, ok := component.Get(ctx)
|
|
if ok {
|
|
pm.handler = handler
|
|
pm.logger = slog.New(handler).With(slog.Any("component", comp))
|
|
}
|
|
}
|
|
|
|
pm.wg.Add(1)
|
|
go pm.handlePoolEvents()
|
|
|
|
return pm
|
|
}
|
|
|
|
func (pm *Postmaster) Send(
|
|
ctx context.Context,
|
|
peerID string,
|
|
data Envelope,
|
|
onOutcome func(LetterOutcome),
|
|
opts ...SendOption,
|
|
) (context.CancelFunc, error) {
|
|
cfg := sendConfig{deadline: pm.cfg.defaultDeadline}
|
|
for _, opt := range opts {
|
|
opt(&cfg)
|
|
}
|
|
|
|
pm.mu.RLock()
|
|
defer pm.mu.RUnlock()
|
|
|
|
// check if peer courier exists
|
|
peerID, ok := pm.poolHasPeer(peerID)
|
|
if !ok {
|
|
return nil, fmt.Errorf("peer not found")
|
|
}
|
|
courier, ok := pm.couriers[peerID]
|
|
if !ok {
|
|
return nil, fmt.Errorf("peer not found")
|
|
}
|
|
|
|
ctx, cancel := context.WithTimeout(ctx, cfg.deadline)
|
|
letter := OutboundLetter{
|
|
id: pm.counter.Add(1),
|
|
peerID: peerID,
|
|
data: data,
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
}
|
|
|
|
courier.Enqueue(letter, onOutcome)
|
|
|
|
return cancel, nil
|
|
}
|
|
|
|
func (pm *Postmaster) Peers() []string {
|
|
pm.mu.RLock()
|
|
defer pm.mu.RUnlock()
|
|
|
|
peers := make([]string, 0, len(pm.couriers))
|
|
for id, _ := range pm.couriers {
|
|
peers = append(peers, id)
|
|
}
|
|
return peers
|
|
}
|
|
|
|
func (pm *Postmaster) Close() {
|
|
pm.cancel()
|
|
pm.wg.Wait()
|
|
|
|
// close each courier
|
|
pm.mu.Lock()
|
|
couriers := pm.couriers
|
|
pm.couriers = make(map[string]*Courier)
|
|
pm.mu.Unlock()
|
|
|
|
for _, courier := range couriers {
|
|
courier.Close()
|
|
}
|
|
}
|
|
|
|
func (pm *Postmaster) handlePoolEvents() {
|
|
defer pm.wg.Done()
|
|
|
|
for {
|
|
select {
|
|
case <-pm.ctx.Done():
|
|
return
|
|
case ev := <-pm.poolEvents:
|
|
switch ev.Kind {
|
|
case EventAdded:
|
|
pm.mu.Lock()
|
|
_, exists := pm.couriers[ev.ID]
|
|
if exists {
|
|
pm.mu.Unlock()
|
|
continue
|
|
}
|
|
send := func(data Envelope) error { return pm.poolSend(ev.ID, data) }
|
|
courier := NewCourier(pm.ctx, send, pm.handler)
|
|
pm.couriers[ev.ID] = courier
|
|
pm.mu.Unlock()
|
|
|
|
case EventRemoved:
|
|
pm.mu.Lock()
|
|
courier, exists := pm.couriers[ev.ID]
|
|
if exists {
|
|
delete(pm.couriers, ev.ID)
|
|
}
|
|
pm.mu.Unlock()
|
|
courier.Close()
|
|
|
|
case EventConnected:
|
|
pm.mu.RLock()
|
|
courier, exists := pm.couriers[ev.ID]
|
|
if exists {
|
|
courier.HandleConnect()
|
|
}
|
|
pm.mu.RUnlock()
|
|
|
|
case EventDisconnected:
|
|
pm.mu.RLock()
|
|
courier, exists := pm.couriers[ev.ID]
|
|
if exists {
|
|
courier.HandleDisconnect()
|
|
}
|
|
pm.mu.RUnlock()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Courier
|
|
// ----------------------------------------------------------------------------
|
|
|
|
// Traveller
|
|
|
|
type letterTraveller struct {
|
|
letter OutboundLetter
|
|
onOutcome func(LetterOutcome)
|
|
|
|
sentAt time.Time
|
|
missedAt time.Time
|
|
retries int
|
|
once sync.Once
|
|
}
|
|
|
|
func (t *letterTraveller) isCancelled() bool {
|
|
return t.letter.ctx.Err() != nil
|
|
}
|
|
|
|
func (t *letterTraveller) countRetry() { t.retries++ }
|
|
func (t *letterTraveller) setSentAt(at time.Time) { t.sentAt = at }
|
|
func (t *letterTraveller) setMissedAt(at time.Time) { t.missedAt = at }
|
|
|
|
// Courier
|
|
|
|
func NewCourier(
|
|
ctx context.Context,
|
|
sendFunc func(data Envelope) error, // func => PoolSendFunc(id)
|
|
handler slog.Handler,
|
|
) *Courier {
|
|
ctx, cancel := context.WithCancel(
|
|
component.MustExtend(ctx, "courier"))
|
|
|
|
c := &Courier{
|
|
cmd: make(chan courierCommand, 64),
|
|
sendFunc: sendFunc,
|
|
ctx: ctx,
|
|
cancel: cancel,
|
|
}
|
|
|
|
if handler != nil {
|
|
comp, ok := component.Get(ctx)
|
|
if ok {
|
|
c.logger = slog.New(handler).With(slog.Any("component", comp))
|
|
}
|
|
}
|
|
|
|
c.wg.Add(1)
|
|
go c.run()
|
|
|
|
return c
|
|
}
|
|
|
|
func (c *Courier) Enqueue(letter OutboundLetter, onOutcome func(LetterOutcome)) {
|
|
traveller := &letterTraveller{
|
|
letter: letter,
|
|
onOutcome: onOutcome,
|
|
}
|
|
c.command(cmdEnqueue{traveller: traveller})
|
|
}
|
|
|
|
func (c *Courier) HandleConnect() {
|
|
c.command(cmdHandleConnect{})
|
|
}
|
|
|
|
func (c *Courier) HandleDisconnect() {
|
|
c.command(cmdHandleDisconnect{})
|
|
}
|
|
|
|
func (c *Courier) Close() {
|
|
c.cancel()
|
|
c.wg.Wait()
|
|
// cancel remaining letters
|
|
for {
|
|
t, ok := c.pop()
|
|
if !ok {
|
|
break
|
|
}
|
|
t.letter.cancel()
|
|
t.setMissedAt(time.Now())
|
|
c.doneOnce(t)
|
|
}
|
|
}
|
|
|
|
// Internal
|
|
|
|
func (c *Courier) command(cmd courierCommand) {
|
|
select {
|
|
case <-c.ctx.Done():
|
|
case c.cmd <- cmd:
|
|
}
|
|
}
|
|
|
|
func (c *Courier) run() {
|
|
defer c.wg.Done()
|
|
|
|
for {
|
|
select {
|
|
case <-c.ctx.Done():
|
|
return
|
|
case cmd := <-c.cmd:
|
|
cmd.apply(c)
|
|
c.maybeSend()
|
|
}
|
|
}
|
|
}
|
|
|
|
func (c *Courier) maybeSend() {
|
|
if !c.preflight() {
|
|
c.drain()
|
|
return
|
|
}
|
|
|
|
t, ok := c.pop()
|
|
if !ok {
|
|
return
|
|
}
|
|
c.sending = true
|
|
|
|
c.wg.Add(1)
|
|
go c.sendOnce(t)
|
|
|
|
}
|
|
|
|
func (c *Courier) sendOnce(t *letterTraveller) {
|
|
defer c.wg.Done()
|
|
err := c.sendFunc(t.letter.data)
|
|
c.command(cmdHandleSendResult{traveller: t, at: time.Now(), err: err})
|
|
}
|
|
|
|
func (c *Courier) doneOnce(t *letterTraveller) {
|
|
var kind LetterOutcomeKind
|
|
if t.isCancelled() {
|
|
// letter was cancelled
|
|
if t.letter.ctx.Err() == context.DeadlineExceeded {
|
|
// letter expired
|
|
kind = OutcomeExpired
|
|
} else {
|
|
// letter was cancelled externally
|
|
kind = OutcomeCancelled
|
|
}
|
|
} else {
|
|
// letter was sent
|
|
kind = OutcomeSent
|
|
}
|
|
|
|
outcome := LetterOutcome{
|
|
LetterID: t.letter.id,
|
|
PeerID: t.letter.peerID,
|
|
Kind: kind,
|
|
SentAt: t.sentAt,
|
|
MissedAt: t.missedAt,
|
|
Retries: t.retries,
|
|
}
|
|
|
|
t.once.Do(func() {
|
|
t.letter.cancel()
|
|
go t.onOutcome(outcome)
|
|
})
|
|
}
|
|
|
|
// Helpers
|
|
|
|
func (c *Courier) preflight() bool {
|
|
isConnected := c.connected
|
|
notAlreadySending := !c.sending
|
|
hasQueuedLetters := c.queue.Len() > 0
|
|
return isConnected && notAlreadySending && hasQueuedLetters
|
|
}
|
|
|
|
func (c *Courier) drain() {
|
|
for {
|
|
front := c.queue.Front()
|
|
if front == nil {
|
|
return
|
|
}
|
|
|
|
t := front.Value.(*letterTraveller)
|
|
if !t.isCancelled() {
|
|
return
|
|
}
|
|
|
|
t.setMissedAt(time.Now())
|
|
c.doneOnce(t)
|
|
c.queue.Remove(front)
|
|
}
|
|
}
|
|
|
|
func (c *Courier) pop() (*letterTraveller, bool) {
|
|
for {
|
|
front := c.queue.Front()
|
|
if front == nil {
|
|
return nil, false
|
|
}
|
|
|
|
t := front.Value.(*letterTraveller)
|
|
c.queue.Remove(front)
|
|
|
|
if !t.isCancelled() {
|
|
return t, true
|
|
}
|
|
|
|
t.setMissedAt(time.Now())
|
|
c.doneOnce(t)
|
|
}
|
|
}
|
|
|
|
// ----------------------------------------------------------------------------
|
|
// Commands
|
|
// ----------------------------------------------------------------------------
|
|
|
|
type cmdEnqueue struct{ traveller *letterTraveller }
|
|
|
|
func (cmd cmdEnqueue) apply(c *Courier) {
|
|
c.queue.PushBack(cmd.traveller)
|
|
}
|
|
|
|
type cmdHandleConnect struct{}
|
|
|
|
func (cmd cmdHandleConnect) apply(c *Courier) {
|
|
c.connected = true
|
|
}
|
|
|
|
type cmdHandleDisconnect struct{}
|
|
|
|
func (cmd cmdHandleDisconnect) apply(c *Courier) {
|
|
c.connected = false
|
|
}
|
|
|
|
type cmdHandleSendResult struct {
|
|
traveller *letterTraveller
|
|
at time.Time
|
|
err error
|
|
}
|
|
|
|
func (cmd cmdHandleSendResult) apply(c *Courier) {
|
|
c.sending = false
|
|
if cmd.err != nil {
|
|
cmd.traveller.countRetry()
|
|
c.queue.PushFront(cmd.traveller)
|
|
} else {
|
|
cmd.traveller.setSentAt(cmd.at)
|
|
c.doneOnce(cmd.traveller)
|
|
}
|
|
}
|