im_server/server/db/mysql/group_message_shards.go

//go:build mysql
// +build mysql

package mysql

import (
	"context"
	"fmt"
	"time"

	"github.com/jmoiron/sqlx"
	"github.com/tinode/chat/server/store"
	t "github.com/tinode/chat/server/store/types"
)

func ensureGroupMessageShardInfrastructure(exec sqlExecutor, tenantID t.TenantID) error {
	statements := []string{
		`CREATE TABLE IF NOT EXISTS im_group_message_shards (
			tenant_id BIGINT NOT NULL DEFAULT 0, shard_month CHAR(6) NOT NULL, table_name VARCHAR(64) NOT NULL,
			shard_start DATETIME(3) NOT NULL, shard_end DATETIME(3) NOT NULL,
			status VARCHAR(16) NOT NULL DEFAULT 'prepared', row_count BIGINT NOT NULL DEFAULT 0,
			min_created_at DATETIME(3), max_created_at DATETIME(3),
			created_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3),
			updated_at DATETIME(3) NOT NULL DEFAULT CURRENT_TIMESTAMP(3) ON UPDATE CURRENT_TIMESTAMP(3),
			PRIMARY KEY(tenant_id, shard_month), UNIQUE INDEX uk_im_group_message_shards_table(tenant_id, table_name),
			INDEX idx_im_group_message_shards_status(tenant_id, status, shard_month)
		) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci`,
	}
	for _, statement := range statements {
		if _, err := exec.Exec(statement); err != nil {
			return err
		}
	}

	return ensureGroupMessageShardWindow(exec, tenantID, time.Now().UTC())
}

func ensureGroupMessageShardWindow(exec sqlExecutor, tenantID t.TenantID, now time.Time) error {
	first := time.Date(now.Year(), now.Month(), 1, 0, 0, 0, 0, time.UTC)
	for _, offset := range []int{-1, 0, 1, 2} {
		start := first.AddDate(0, offset, 0)
		status := "prepared"
		if offset == -1 {
			status = "sealing"
		} else if offset == 0 {
			status = "writable"
		}
		if err := ensureGroupMessageShard(exec, tenantID, start, status); err != nil {
			return err
		}
	}
	_, err := exec.Exec("UPDATE im_group_message_shards SET status='sealed' WHERE tenant_id=? AND status='sealing' AND shard_end<?", tenantID, first.AddDate(0, 0, -7))
	return err
}

func (a *adapter) ensureCurrentGroupMessageShardWindow() error {
	month := time.Now().UTC().Format("200601")
	a.tenantMu.Lock()
	defer a.tenantMu.Unlock()
	if a.groupShardWindow == month {
		return nil
	}
	if err := ensureGroupMessageShardWindow(a.db, a.tenantID, time.Now().UTC()); err != nil {
		return err
	}
	a.groupShardWindow = month
	return nil
}

func ensureGroupMessageShard(exec sqlExecutor, tenantID t.TenantID, start time.Time, status string) error {
	month := start.UTC().Format("200601")
	table, ok := tenantGroupMessagesTable(tenantID, month)
	if !ok {
		return fmt.Errorf("mysql: invalid group message shard month %q", month)
	}
	ddl := fmt.Sprintf(`CREATE TABLE IF NOT EXISTS %s (
		tenant_id BIGINT NOT NULL DEFAULT 0, message_id BIGINT NOT NULL, topic VARCHAR(64) NOT NULL, group_id BIGINT NOT NULL,
		seqid BIGINT NOT NULL, from_user BIGINT NOT NULL, client_msg_id VARCHAR(64),
		msg_type VARCHAR(32) NOT NULL DEFAULT 'text', status VARCHAR(32) NOT NULL DEFAULT 'normal',
		content JSON, content_nonce VARBINARY(12), key_version INT NOT NULL DEFAULT 0,
		head JSON, mentioned_user_ids JSON, mention_all TINYINT(1) NOT NULL DEFAULT 0,
		reply_seqid BIGINT, root_seqid BIGINT, file_count INT NOT NULL DEFAULT 0,
		content_size INT NOT NULL DEFAULT 0, has_risk TINYINT(1) NOT NULL DEFAULT 0,
		audit_status VARCHAR(32) NOT NULL DEFAULT 'pass', version INT NOT NULL DEFAULT 1,
		recalled_at DATETIME(3), recalled_by BIGINT, deleted_at DATETIME(3), expires_at DATETIME(3),
		created_at DATETIME(3) NOT NULL, updated_at DATETIME(3) NOT NULL,
		PRIMARY KEY(message_id), UNIQUE INDEX uk_%s_group_seq(group_id, seqid),
		UNIQUE INDEX uk_%s_topic_seq(topic, seqid),
		UNIQUE INDEX uk_%s_sender_client(topic, from_user, client_msg_id),
		INDEX idx_%s_group_status_seq(group_id, status, seqid),
		INDEX idx_%s_topic_created(topic, created_at), INDEX idx_%s_topic_from_seq(topic, from_user, seqid),
		INDEX idx_%s_audit_created(audit_status, created_at), INDEX idx_%s_expires(expires_at)
	) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_0900_ai_ci`,
		table, month, month, month, month, month, month, month, month)
	if _, err := exec.Exec(ddl); err != nil {
		return err
	}
	end := start.AddDate(0, 1, 0)
	_, err := exec.Exec(`INSERT INTO im_group_message_shards
		(tenant_id,shard_month,table_name,shard_start,shard_end,status)
		VALUES(?,?,?,?,?,?) ON DUPLICATE KEY UPDATE tenant_id=VALUES(tenant_id),table_name=VALUES(table_name),
		shard_start=VALUES(shard_start),shard_end=VALUES(shard_end),
		status=IF(status IN ('sealed','archived'),status,VALUES(status))`,
		tenantID, month, table, start, end, status)
	return err
}

func groupMessageShardTable(tenantID t.TenantID, month string) (string, bool) {
	return tenantGroupMessagesTable(tenantID, month)
}

func groupMessageTablesForTopic(q sqlx.QueryerContext, ctx context.Context, tenantID t.TenantID, topic string) ([]string, error) {
	routes, cached, err := store.GetGroupMessageShardRoutes(tenantID, topic)
	if err != nil {
		cached = false
	}
	if !cached {
		if routes, err = groupMessageAllShardRoutes(q, ctx, tenantID); err != nil {
			return nil, err
		}
	}
	tables := make([]string, 0, len(routes)+1)
	for _, route := range routes {
		if table, ok := groupMessageShardTable(tenantID, route.Month); ok {
			tables = append(tables, table)
		}
	}
	return tables, nil
}

func groupMessageTablesForSeq(q sqlx.QueryerContext, ctx context.Context, tenantID t.TenantID, topic string, seq int64) ([]string, error) {
	routes, cached, err := store.GetGroupMessageShardRoutes(tenantID, topic)
	if err != nil {
		cached = false
	}
	if !cached {
		if routes, err = groupMessageAllShardRoutes(q, ctx, tenantID); err != nil {
			return nil, err
		}
	}
	tables := make([]string, 0, 2)
	for _, route := range routes {
		if route.MinSeq <= seq && route.MaxSeq >= seq {
			if table, ok := groupMessageShardTable(tenantID, route.Month); ok {
				tables = append(tables, table)
			}
		}
	}
	return tables, nil
}

func groupMessageAllShardRoutes(q sqlx.QueryerContext, ctx context.Context, tenantID t.TenantID) ([]store.GroupMessageShardRoute, error) {
	var rows []struct {
		Month string `db:"month"`
	}
	if err := sqlx.SelectContext(ctx, q, &rows, `SELECT shard_month AS month FROM im_group_message_shards WHERE tenant_id=? ORDER BY shard_month DESC`, tenantID); err != nil {
		return nil, err
	}
	routes := make([]store.GroupMessageShardRoute, 0, len(rows))
	for _, row := range rows {
		if validShardMonth(row.Month) {
			routes = append(routes, store.GroupMessageShardRoute{Month: row.Month, MinSeq: 1, MaxSeq: 1<<62 - 1, Count: 0})
		}
	}
	return routes, nil
}

func groupMessageMonthFromTable(table string) string {
	const suffixLength = 6
	if len(table) > suffixLength && validShardMonth(table[len(table)-suffixLength:]) {
		return table[len(table)-suffixLength:]
	}
	return ""
}