nixevol/NixMsg
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 仅向已握手连接推送并拆分调度循环

Browse Source
This commit is contained in:
Nixevol
2026-09-30 16:22:43 +08:00
parent ad4f13193c
commit 83521f4b87
14 changed files with 1192 additions and 290 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/app/message/push.go
+311 -171
View File
@@ -4,28 +4,75 @@ import (
"context"
"database/sql"
"encoding/json"
"errors"
"fmt"
"log/slog"
"sync"
"time"
"git.asio.asia/nixevol/NixMsg/internal/app/port"
"git.asio.asia/nixevol/NixMsg/internal/protocol"
)
type dueMsg struct {
seq int64
senderID string
destKind string
destID string
sendAt int64
keep int
ttl int64
receipt int
}
type pushItem struct {
seq int64
sendAt int64
keep int
expireAt sql.NullInt64
msgID string
senderID string
destKind string
destID string
meta string
contentType string
bodyEnc string
body []byte
payload []byte
}
// DispatchDue 分发已到点的 scheduled 消息（按 send_at、seq）。
// limit<=0 时循环到取空或用完时间预算，单条失败只记日志并跳过。
func (a *App) DispatchDue(ctx context.Context, nowMs int64, limit int) (int, error) {
if limit <= 0 {
limit = 64
budgeted := limit <= 0
batch := limit
if batch <= 0 {
batch = 64
}
type due struct {
seq int64
senderID string
destKind string
destID string
sendAt int64
keep int
ttl int64
receipt int
deadline := time.Now().Add(time.Hour)
if budgeted {
deadline = time.Now().Add(dispatchBudget)
}
total := 0
for {
if err := ctx.Err(); err != nil {
return total, err
}
if budgeted && time.Now().After(deadline) {
return total, nil
}
n, err := a.dispatchDueBatch(ctx, nowMs, batch)
total += n
if err != nil {
return total, err
}
if n == 0 || !budgeted {
return total, nil
}
}
}
func (a *App) dispatchDueBatch(ctx context.Context, nowMs int64, limit int) (int, error) {
rows, err := a.db.Read.QueryContext(ctx, `
SELECT seq, sender_id, dest_kind, dest_id, send_at, keep, ttl_seconds, receipt
FROM messages
@@ -35,9 +82,9 @@ LIMIT ?`, nowMs, limit)
if err != nil {
return 0, err
}
var list []due
var list []dueMsg
for rows.Next() {
var d due
var d dueMsg
if err := rows.Scan(&d.seq, &d.senderID, &d.destKind, &d.destID, &d.sendAt, &d.keep, &d.ttl, &d.receipt); err != nil {
_ = rows.Close()
return 0, err
@@ -49,55 +96,68 @@ LIMIT ?`, nowMs, limit)
return 0, err
}
_ = rows.Close()
if len(list) == 0 {
return 0, nil
}
sem := make(chan struct{}, dispatchConcurrency)
var wg sync.WaitGroup
var mu sync.Mutex
n := 0
wake := map[string]struct{}{}
for _, d := range list {
var claimed bool
err := a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
_, c, err := a.dispatchFullTx(tx, d.seq, d.senderID, d.destKind, d.destID, d.sendAt, d.keep, d.ttl, d.receipt != 0, nowMs)
claimed = c
return err
})
if err != nil {
return n, err
}
if !claimed {
continue
}
n++
rows2, qErr := a.db.Read.QueryContext(ctx, `
d := d
wg.Add(1)
sem <- struct{}{}
go func() {
defer wg.Done()
defer func() { <-sem }()
var claimed bool
err := a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
_, c, err := a.dispatchFullTx(tx, d.seq, d.senderID, d.destKind, d.destID, d.sendAt, d.keep, d.ttl, d.receipt != 0, nowMs)
claimed = c
return err
})
if err != nil {
slog.Error("dispatch due item", "seq", d.seq, "err", err)
return
}
if !claimed {
return
}
mu.Lock()
n++
mu.Unlock()
rows2, qErr := a.db.Read.QueryContext(ctx, `
SELECT DISTINCT endpoint_id FROM deliveries WHERE seq = ? AND state = 'pending'`, d.seq)
if qErr == nil {
if qErr != nil {
return
}
for rows2.Next() {
var ep string
if rows2.Scan(&ep) == nil {
mu.Lock()
wake[ep] = struct{}{}
mu.Unlock()
}
}
_ = rows2.Close()
}
}()
}
wg.Wait()
for ep := range wake {
a.WakePush(ep)
}
return n, nil
}
// PushPending 向指定连接推送 pending 投递与回执。
// PushPending 向已握手连接推送 pending 投递与回执。
func (a *App) PushPending(ctx context.Context, endpointID string, connID port.ConnID) error {
nowMs := a.now().UnixMilli()
live, ok := a.lookupConn(endpointID)
if !ok || (connID != "" && live.ConnID != connID) {
// 仍处理该代号上的确认超时与清标记场景：用传入 connID
if connID == "" {
return nil
}
live = LiveConn{ConnID: connID}
} else if connID == "" {
connID = live.ConnID
live, connID, ok := a.canPush(endpointID, connID)
if !ok {
return nil
}
nowMs := a.now().UnixMilli()
if err := a.processAckTimeouts(ctx, endpointID, connID, nowMs); err != nil {
return err
}
@@ -128,31 +188,13 @@ LIMIT ?`, endpointID, room)
if err != nil {
return err
}
type item struct {
seq int64
sendAt int64
keep int
expireAt sql.NullInt64
msgID string
senderID string
destKind string
destID string
meta string
contentType string
bodyEnc string
body []byte
}
var items []item
var items []pushItem
for rows.Next() {
var it item
var body sql.NullString
var bodyBlob []byte
if err := rows.Scan(&it.seq, &it.sendAt, &it.keep, &it.expireAt, &it.msgID, &it.senderID, &it.destKind, &it.destID, &it.meta, &it.contentType, &it.bodyEnc, &bodyBlob); err != nil {
var it pushItem
if err := rows.Scan(&it.seq, &it.sendAt, &it.keep, &it.expireAt, &it.msgID, &it.senderID, &it.destKind, &it.destID, &it.meta, &it.contentType, &it.bodyEnc, &it.body); err != nil {
_ = rows.Close()
return err
}
_ = body
it.body = bodyBlob
items = append(items, it)
}
_ = rows.Close()
@@ -160,9 +202,9 @@ LIMIT ?`, endpointID, room)
return err
}
var toClaim []pushItem
for _, it := range items {
if it.body == nil {
// 正文已删则跳过（异常）
continue
}
msg := protocol.Msg{
@@ -186,9 +228,40 @@ LIMIT ?`, endpointID, room)
}
continue
}
it.payload = payload
toClaim = append(toClaim, it)
}
claimed := false
err = a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
claimed, err := a.claimPushBatch(ctx, endpointID, connID, nowMs, toClaim)
if err != nil {
return err
}
for _, it := range claimed {
if a.down == nil {
continue
}
pubErr := a.down.PublishDown(ctx, endpointID, connID, it.payload, port.PublishOpts{QoS: 1})
if pubErr != nil {
short, cancel := shortWriteCtx()
_ = a.clearPushed(short, it.seq, endpointID, connID, a.now().UnixMilli())
cancel()
a.scheduleRepush(endpointID, time.Second)
} else {
a.observeDispatchToPush(it.sendAt, nowMs)
}
}
return a.pushReceipts(ctx, endpointID, connID, nowMs)
}
func (a *App) claimPushBatch(ctx context.Context, endpointID string, connID port.ConnID, nowMs int64, items []pushItem) ([]pushItem, error) {
if len(items) == 0 {
return nil, nil
}
var claimed []pushItem
err := a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
claimed = claimed[:0]
for _, it := range items {
res, e := tx.Exec(`
UPDATE deliveries SET pushed_conn = ?, pushed_at = ?, attempts = attempts + 1, updated_at = ?
WHERE seq = ? AND endpoint_id = ? AND state = 'pending' AND pushed_conn IS NULL`,
@@ -197,43 +270,35 @@ WHERE seq = ? AND endpoint_id = ? AND state = 'pending' AND pushed_conn IS NULL`
return e
}
aff, _ := res.RowsAffected()
claimed = aff > 0
return nil
})
if err != nil {
return err
}
if !claimed {
continue
}
large := len(payload) > largeFrameBytes
if large {
if !a.acquireLarge(ctx) {
_ = a.clearPushed(ctx, it.seq, endpointID, connID, nowMs)
continue
if aff > 0 {
claimed = append(claimed, it)
}
a.trackLarge(it.seq, endpointID, true)
}
if a.down == nil {
if large {
a.releaseLarge(it.seq, endpointID)
}
continue
}
pubErr := a.down.PublishDown(ctx, endpointID, connID, payload, port.PublishOpts{QoS: 1})
if pubErr != nil {
_ = a.clearPushed(ctx, it.seq, endpointID, connID, nowMs)
if large {
a.releaseLarge(it.seq, endpointID)
}
return nil
})
if err != nil {
if ctx.Err() != nil {
short, cancel := shortWriteCtx()
a.clearIfClaimed(short, items, endpointID, connID)
cancel()
a.scheduleRepush(endpointID, time.Second)
} else {
a.observeDispatchToPush(it.sendAt, nowMs)
}
return nil, err
}
return claimed, nil
}
func (a *App) clearIfClaimed(ctx context.Context, items []pushItem, endpointID string, connID port.ConnID) {
nowMs := a.now().UnixMilli()
for _, it := range items {
var pushed sql.NullString
_ = a.db.Read.QueryRowContext(ctx, `
SELECT pushed_conn FROM deliveries WHERE seq = ? AND endpoint_id = ? AND state = 'pending'`,
it.seq, endpointID).Scan(&pushed)
if pushed.Valid && pushed.String == string(connID) {
_ = a.clearPushed(ctx, it.seq, endpointID, connID, nowMs)
}
}
return a.pushReceipts(ctx, endpointID, connID, nowMs)
}
func (a *App) observeDispatchToPush(sendAtMs, pushedAtMs int64) {
@@ -256,7 +321,7 @@ WHERE seq = ? AND endpoint_id = ? AND state = 'pending' AND pushed_conn IS NULL`
if aff == 0 {
return nil
}
if err := insertReceiptTx(tx, senderID, seq, endpointID, DeliveryRejected, ReasonTooLarge, nowMs); err != nil {
if _, err := insertReceiptTx(tx, senderID, seq, endpointID, DeliveryRejected, ReasonTooLarge, nowMs); err != nil {
return err
}
return TryFinalizeTx(tx, seq, nowMs, a.lim.RecordRetentionDays)
@@ -264,6 +329,11 @@ WHERE seq = ? AND endpoint_id = ? AND state = 'pending' AND pushed_conn IS NULL`
}
func (a *App) clearPushed(ctx context.Context, seq int64, endpointID string, connID port.ConnID, nowMs int64) error {
if err := ctx.Err(); err != nil {
var cancel context.CancelFunc
ctx, cancel = shortWriteCtx()
defer cancel()
}
return a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
_, err := tx.Exec(`
UPDATE deliveries SET pushed_conn = NULL, updated_at = ?
@@ -291,14 +361,14 @@ WHERE d.endpoint_id = ? AND d.state = 'pending' AND d.pushed_conn = ?
seq int64
keep int
expireAt sql.NullInt64
pushedAt int64
senderID string
msgID string
}
var list []to
for rows.Next() {
var t to
var pushedAt int64
if err := rows.Scan(&t.seq, &t.keep, &t.expireAt, &pushedAt, &t.senderID, &t.msgID); err != nil {
if err := rows.Scan(&t.seq, &t.keep, &t.expireAt, &t.pushedAt, &t.senderID, &t.msgID); err != nil {
_ = rows.Close()
return err
}
@@ -311,11 +381,12 @@ WHERE d.endpoint_id = ? AND d.state = 'pending' AND d.pushed_conn = ?
var keep int
var expireAt sql.NullInt64
var pushedConn sql.NullString
var pushedAt sql.NullInt64
err := tx.QueryRow(`
SELECT keep, expire_at, pushed_conn FROM deliveries
WHERE seq = ? AND endpoint_id = ? AND state = 'pending'`, t.seq, endpointID).Scan(&keep, &expireAt, &pushedConn)
SELECT keep, expire_at, pushed_conn, pushed_at FROM deliveries
WHERE seq = ? AND endpoint_id = ? AND state = 'pending'`, t.seq, endpointID).Scan(&keep, &expireAt, &pushedConn, &pushedAt)
if err != nil {
if err == sql.ErrNoRows {
if errors.Is(err, sql.ErrNoRows) {
return nil
}
return err
@@ -323,13 +394,15 @@ WHERE seq = ? AND endpoint_id = ? AND state = 'pending'`, t.seq, endpointID).Sca
if !pushedConn.Valid || pushedConn.String != string(connID) {
return nil
}
if !pushedAt.Valid || pushedAt.Int64 > nowMs-timeoutMs {
return nil
}
if keep == 0 {
return a.finishDeliveryTx(tx, t.seq, endpointID, t.senderID, t.msgID, DeliveryDropped, ReasonNotAcked, true, nowMs)
}
if expireAt.Valid && expireAt.Int64 <= nowMs {
return a.finishDeliveryTx(tx, t.seq, endpointID, t.senderID, t.msgID, DeliveryExpired, ReasonTTL, true, nowMs)
}
// 重推：清标记
_, err = tx.Exec(`
UPDATE deliveries SET pushed_conn = NULL, updated_at = ?
WHERE seq = ? AND endpoint_id = ? AND state = 'pending' AND pushed_conn = ?`,
@@ -339,7 +412,6 @@ WHERE seq = ? AND endpoint_id = ? AND state = 'pending' AND pushed_conn = ?`,
if err != nil {
return err
}
a.releaseLarge(t.seq, endpointID)
}
return nil
}
@@ -357,7 +429,7 @@ WHERE seq = ? AND endpoint_id = ? AND state = 'pending'`,
return nil
}
if state != DeliveryRecalled {
if err := insertReceiptTx(tx, senderID, seq, endpointID, state, reason, nowMs); err != nil {
if _, err := insertReceiptTx(tx, senderID, seq, endpointID, state, reason, nowMs); err != nil {
return err
}
}
@@ -406,6 +478,7 @@ func (a *App) flushRevokes(ctx context.Context) {
if a.down == nil {
return
}
var retry []revokeJob
for _, j := range jobs {
frame := protocol.Revoked{
V: protocol.Version, Type: protocol.TypeRevoked,
@@ -415,26 +488,48 @@ func (a *App) flushRevokes(ctx context.Context) {
if err != nil {
continue
}
_ = a.down.PublishDown(ctx, j.endpointID, j.connID, payload, port.PublishOpts{QoS: 1})
if err := a.down.PublishDown(ctx, j.endpointID, j.connID, payload, port.PublishOpts{QoS: 1}); err != nil {
retry = append(retry, j)
a.scheduleRepush(j.endpointID, time.Second)
}
}
if len(retry) > 0 {
a.mu.Lock()
a.pendingRevoke = append(retry, a.pendingRevoke...)
a.mu.Unlock()
}
}
// OnPublishDropped 清推送标记并 1 秒后重推。
func (a *App) OnPublishDropped(ctx context.Context, endpointID string, connID port.ConnID, payload []byte) error {
var head struct {
Type string `json:"type"`
ID string `json:"id"`
From string `json:"from"`
Type string `json:"type"`
ID string `json:"id"`
From string `json:"from"`
ReceiptID string `json:"receipt_id"`
}
if err := json.Unmarshal(payload, &head); err != nil || head.Type != protocol.TypeMsg {
if err := json.Unmarshal(payload, &head); err != nil {
return nil
}
switch head.Type {
case protocol.TypeReceipt:
if rid, err := parseReceiptID(head.ReceiptID); err == nil {
a.unmarkReceipt(string(connID), rid)
}
a.scheduleRepush(endpointID, time.Second)
return nil
case protocol.TypeMsg:
default:
return nil
}
nowMs := a.now().UnixMilli()
err := a.db.Queue.Do(ctx, func(tx *sql.Tx) error {
short, cancel := shortWriteCtx()
defer cancel()
err := a.db.Queue.Do(short, func(tx *sql.Tx) error {
var seq int64
err := tx.QueryRow(`SELECT seq FROM messages WHERE sender_id = ? AND id = ?`, head.From, head.ID).Scan(&seq)
if err != nil {
if err == sql.ErrNoRows {
if errors.Is(err, sql.ErrNoRows) {
return nil
}
return err
@@ -443,7 +538,6 @@ func (a *App) OnPublishDropped(ctx context.Context, endpointID string, connID po
UPDATE deliveries SET pushed_conn = NULL, updated_at = ?
WHERE seq = ? AND endpoint_id = ? AND state = 'pending' AND pushed_conn = ?`,
nowMs, seq, endpointID, string(connID))
a.releaseLarge(seq, endpointID)
return err
})
if err != nil {
@@ -458,77 +552,128 @@ func (a *App) pushReceipts(ctx context.Context, endpointID string, connID port.C
if window <= 0 {
window = defaultReceiptWindow
}
// 简化：未单独记 inflight 回执，按未确认回执取窗口条数
retryAfter := a.receiptRetryMs()
a.mu.Lock()
held := a.rcptInflight[string(connID)]
inflightN := 0
stale := map[int64]struct{}{}
for rid, at := range held {
if nowMs-at >= retryAfter {
stale[rid] = struct{}{}
continue
}
inflightN++
}
a.mu.Unlock()
room := window - inflightN
if room <= 0 {
return nil
}
rows, err := a.db.Read.QueryContext(ctx, `
SELECT receipt_id, msg_id, endpoint_id, state, reason, created_at
FROM receipts
WHERE sender_id = ? AND acked = 0
ORDER BY receipt_id ASC
LIMIT ?`, endpointID, window)
LIMIT ?`, endpointID, window+len(stale))
if err != nil {
return err
}
defer func() { _ = rows.Close() }()
type rcpt struct {
rid int64
msgID, epID string
state, reason string
created int64
}
var list []rcpt
for rows.Next() {
var r rcpt
if err := rows.Scan(&r.rid, &r.msgID, &r.epID, &r.state, &r.reason, &r.created); err != nil {
_ = rows.Close()
return err
}
list = append(list, r)
}
if err := rows.Err(); err != nil {
_ = rows.Close()
return err
}
_ = rows.Close()
if a.down == nil {
return nil
}
for rows.Next() {
var rid int64
var msgID, epID, state, reason string
var created int64
if err := rows.Scan(&rid, &msgID, &epID, &state, &reason, &created); err != nil {
return err
sent := 0
for _, r := range list {
if sent >= room {
break
}
a.mu.Lock()
m := a.rcptInflight[string(connID)]
at, in := m[r.rid]
fresh := in && nowMs-at < retryAfter
a.mu.Unlock()
if fresh {
continue
}
a.markReceipt(string(connID), r.rid, nowMs)
frame := protocol.Receipt{
V: protocol.Version, Type: protocol.TypeReceipt,
ReceiptID: fmt.Sprintf("%d", rid), ID: msgID, EndpointID: epID,
State: state, Reason: reason, AtMs: created,
ReceiptID: fmt.Sprintf("%d", r.rid), ID: r.msgID, EndpointID: r.epID,
State: r.state, Reason: r.reason, AtMs: r.created,
}
payload, err := protocol.Marshal(frame)
if err != nil {
a.unmarkReceipt(string(connID), r.rid)
return err
}
_ = a.down.PublishDown(ctx, endpointID, connID, payload, port.PublishOpts{QoS: 1})
}
_ = nowMs
return rows.Err()
}
func (a *App) acquireLarge(ctx context.Context) bool {
select {
case a.largeSem <- struct{}{}:
return true
case <-ctx.Done():
return false
default:
return false
}
}
func (a *App) trackLarge(seq int64, endpointID string, hold bool) {
a.mu.Lock()
defer a.mu.Unlock()
key := largeKey(seq, endpointID)
if hold {
a.largeHeld[key] = true
}
}
func (a *App) releaseLarge(seq int64, endpointID string) {
a.mu.Lock()
defer a.mu.Unlock()
key := largeKey(seq, endpointID)
if a.largeHeld[key] {
delete(a.largeHeld, key)
select {
case <-a.largeSem:
default:
if err := a.down.PublishDown(ctx, endpointID, connID, payload, port.PublishOpts{QoS: 1}); err != nil {
a.unmarkReceipt(string(connID), r.rid)
a.scheduleRepush(endpointID, time.Second)
continue
}
sent++
}
return nil
}
func (a *App) receiptRetryMs() int64 {
sec := a.lim.AckTimeoutSeconds
if sec <= 0 {
sec = 60
}
return sec * 1000
}
func (a *App) markReceipt(connID string, rid, nowMs int64) {
a.mu.Lock()
defer a.mu.Unlock()
m := a.rcptInflight[connID]
if m == nil {
m = make(map[int64]int64)
a.rcptInflight[connID] = m
}
m[rid] = nowMs
}
func (a *App) unmarkReceipt(connID string, rid int64) {
a.mu.Lock()
defer a.mu.Unlock()
if m := a.rcptInflight[connID]; m != nil {
delete(m, rid)
}
}
func largeKey(seq int64, endpointID string) string {
return fmt.Sprintf("%d:%s", seq, endpointID)
func (a *App) clearReceiptInflight(connID string) {
a.mu.Lock()
defer a.mu.Unlock()
delete(a.rcptInflight, connID)
}
func parseReceiptID(s string) (int64, error) {
var n int64
_, err := fmt.Sscan(s, &n)
return n, err
}
func (a *App) scheduleRepush(endpointID string, d time.Duration) {
@@ -545,16 +690,11 @@ func (a *App) scheduleRepush(endpointID string, d time.Duration) {
})
}
// WakePush 唤醒推送；若有登记的连接则异步 PushPending。
// WakePush 唤醒该端已握手连接的推送 worker（合并唤醒）。
func (a *App) WakePush(endpointID string) {
live, ok := a.lookupConn(endpointID)
if !ok || a.down == nil {
_, _, ok := a.canPush(endpointID, "")
if !ok {
return
}
go func() {
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
defer cancel()
_ = a.PushPending(ctx, endpointID, live.ConnID)
a.flushRevokes(ctx)
}()
a.signalWorker(endpointID)
}