fix: harden high-fanout delivery under load

server/db/mysql/adapter.go

@@ -4836,6 +4836,27 @@ func isDupe(err error) bool {
 	return ok && myerr.Number == 1062
 }
 
+func isInvalidConnection(err error) bool {
+	return errors.Is(err, ms.ErrInvalidConn)
+}
+
+// retryInvalidConnection retries only operations which are independently
+// idempotent. Callers must not use it for unconstrained writes.
+func retryInvalidConnection(operation string, fn func() error) error {
+	err := fn()
+	if !isInvalidConnection(err) {
+		return err
+	}
+	immetrics.IncDBInvalidConnectionRetry(operation, "attempt")
+	err = fn()
+	if err == nil || isDupe(err) {
+		immetrics.IncDBInvalidConnectionRetry(operation, "recovered")
+	} else {
+		immetrics.IncDBInvalidConnectionRetry(operation, "failed")
+	}
+	return err
+}
+
 func isMissingTable(err error) bool {
 	if err == nil {
 		return false

server/db/mysql/invalid_connection_test.go

@@ -0,0 +1,55 @@
+//go:build mysql
+
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+	"errors"
+	"testing"
+
+	ms "github.com/go-sql-driver/mysql"
+)
+
+func TestRetryInvalidConnectionRetriesOnce(t *testing.T) {
+	attempts := 0
+	err := retryInvalidConnection("test", func() error {
+		attempts++
+		if attempts == 1 {
+			return ms.ErrInvalidConn
+		}
+		return nil
+	})
+	if err != nil || attempts != 2 {
+		t.Fatalf("err=%v attempts=%d", err, attempts)
+	}
+}
+
+func TestRetryInvalidConnectionDoesNotRetryOtherErrors(t *testing.T) {
+	want := errors.New("write failed")
+	attempts := 0
+	err := retryInvalidConnection("test", func() error {
+		attempts++
+		return want
+	})
+	if !errors.Is(err, want) || attempts != 1 {
+		t.Fatalf("err=%v attempts=%d", err, attempts)
+	}
+}
+
+func TestInstrumentedQueryDiscardsInvalidConnection(t *testing.T) {
+	conn := instrumentedConn{conn: invalidQueryConn{}}
+	_, err := conn.QueryContext(context.Background(), "SELECT 1", nil)
+	if !errors.Is(err, driver.ErrBadConn) {
+		t.Fatalf("err=%v want=%v", err, driver.ErrBadConn)
+	}
+}
+
+type invalidQueryConn struct{}
+
+func (invalidQueryConn) Prepare(string) (driver.Stmt, error) { return nil, driver.ErrSkip }
+func (invalidQueryConn) Close() error                        { return nil }
+func (invalidQueryConn) Begin() (driver.Tx, error)           { return nil, driver.ErrSkip }
+func (invalidQueryConn) QueryContext(context.Context, string, []driver.NamedValue) (driver.Rows, error) {
+	return nil, ms.ErrInvalidConn
+}

server/db/mysql/metrics_driver.go

@@ -113,6 +113,12 @@ func (c instrumentedConn) QueryContext(ctx context.Context, query string, args [
 	started := time.Now()
 	rows, err := queryer.QueryContext(ctx, query, args)
 	immetrics.ObserveDB("mysql", "query", started, err)
+	if isInvalidConnection(err) {
+		// Queries are idempotent. Returning ErrBadConn makes database/sql discard
+		// this connection and transparently retry on a fresh one.
+		immetrics.IncDBInvalidConnectionRetry("query", "discard")
+		return nil, driver.ErrBadConn
+	}
 	return rows, err
 }
 
@@ -211,6 +217,10 @@ func (s instrumentedStmt) QueryContext(ctx context.Context, args []driver.NamedV
 	started := time.Now()
 	rows, err := queryer.QueryContext(ctx, args)
 	immetrics.ObserveDB("mysql", "query", started, err)
+	if isInvalidConnection(err) {
+		immetrics.IncDBInvalidConnectionRetry("prepared_query", "discard")
+		return nil, driver.ErrBadConn
+	}
 	return rows, err
 }
 

server/db/mysql/tenant_adapter.go

@@ -2233,21 +2233,29 @@ func (a *adapter) TenantMessageSave(id t.TenantID, msg *t.Message) error {
 	if a.tenantScopedDB {
 		return a.MessageSave(msg)
 	}
-	ctx, cancel := a.getContext()
-	if cancel != nil {
-		defer cancel()
+	save := func() error {
+		ctx, cancel := a.getContext()
+		if cancel != nil {
+			defer cancel()
+		}
+		clientMsgID := sql.NullString{String: msg.ClientMsgID, Valid: msg.ClientMsgID != ""}
+		res, saveErr := a.db.ExecContext(ctx,
+			fmt.Sprintf("INSERT INTO %s(tenant_id,createdat,updatedat,seqid,topic,`from`,client_msg_id,version,head,content,content_nonce,key_version) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)", tenantMessagesTable(id)),
+			id, msg.CreatedAt, msg.UpdatedAt, msg.SeqId, msg.Topic, decodeUID(t.ParseUid(msg.From)), clientMsgID, messageVersion(msg.Version), msg.Head, common.ToJSON(msg.Content), msg.ContentNonce, msg.KeyVersion)
+		if saveErr == nil {
+			lastID, _ := res.LastInsertId()
+			msg.SetUid(t.Uid(lastID))
+		}
+		return saveErr
+	}
+	if msg.ClientMsgID == "" {
+		err = save()
+	} else {
+		err = retryInvalidConnection("tenant_message_save", save)
 	}
-	clientMsgID := sql.NullString{String: msg.ClientMsgID, Valid: msg.ClientMsgID != ""}
-	res, err := a.db.ExecContext(ctx,
-		fmt.Sprintf("INSERT INTO %s(tenant_id,createdat,updatedat,seqid,topic,`from`,client_msg_id,version,head,content,content_nonce,key_version) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)", tenantMessagesTable(id)),
-		id, msg.CreatedAt, msg.UpdatedAt, msg.SeqId, msg.Topic, decodeUID(t.ParseUid(msg.From)), clientMsgID, messageVersion(msg.Version), msg.Head, common.ToJSON(msg.Content), msg.ContentNonce, msg.KeyVersion)
 	if isDupe(err) {
 		return t.ErrDuplicate
 	}
-	if err == nil {
-		lastID, _ := res.LastInsertId()
-		msg.SetUid(t.Uid(lastID))
-	}
 	return err
 }
 
@@ -2349,7 +2357,17 @@ func (a *adapter) TenantGroupMessageSave(id t.TenantID, msg *t.Message) error {
 	if tenantAdp != a {
 		return tenantAdp.TenantGroupMessageSave(id, msg)
 	}
-	if err = a.ensureCurrentGroupMessageShardWindow(); err != nil {
+	err = retryInvalidConnection("tenant_group_message_save", func() error {
+		return a.tenantGroupMessageSaveOnce(id, msg)
+	})
+	if isDupe(err) {
+		return t.ErrDuplicate
+	}
+	return err
+}
+
+func (a *adapter) tenantGroupMessageSaveOnce(id t.TenantID, msg *t.Message) error {
+	if err := a.ensureCurrentGroupMessageShardWindow(); err != nil {
 		return err
 	}
 	ctx, cancel := a.getContext()
@@ -2398,9 +2416,6 @@ func (a *adapter) TenantGroupMessageSave(id t.TenantID, msg *t.Message) error {
 		id, messageID, msg.Topic, groupID, msg.SeqId, decodeUID(t.ParseUid(msg.From)), clientMsgID, messageHeadType(msg.Head), "normal",
 		contentJSON, msg.Head, messageHeadJSON(msg.Head, "mentioned_user_ids"), messageHeadBool(msg.Head, "mention_all"),
 		0, len(contentJSON), messageVersion(msg.Version), msg.CreatedAt, msg.UpdatedAt, msg.ContentNonce, msg.KeyVersion)
-	if isDupe(err) {
-		return t.ErrDuplicate
-	}
 	if err != nil {
 		return err
 	}