yx-dataset-server/library/dingtalk/token.go

package dingtalk

import (
	"context"
	"sync"
	"time"
)

// 钉钉 access_token 缓存相关常量
const (
	// tokenRefreshBuffer token 过期前提前刷新的时间余量
	tokenRefreshBuffer = 2 * time.Minute
	// tokenMaxTTL 兜底 TTL（钉钉 token 有效期一般 7200s）
	tokenMaxTTL = 2 * time.Hour
)

// TokenStorage 可插拔的外部 Token 存储（如 Redis），接口保持最小
// 实现需保证并发安全。返回的 token 为空串视为"未命中/已过期"。
type TokenStorage interface {
	GetToken(ctx context.Context, clientID string) (token string, err error)
	SetToken(ctx context.Context, clientID string, token string, ttl time.Duration) error
}

// tokenEntry 内存 L1 缓存单元
type tokenEntry struct {
	token     string
	expiresAt time.Time
}

// tokenManager Token 管理器：L1 内存 + L2 可插拔（如 Redis），并带 singleflight 去重
type tokenManager struct {
	mu         sync.RWMutex
	entries    map[string]tokenEntry // key: clientID
	inflight   map[string]*inflightCall
	inflightMu sync.Mutex
	storage    TokenStorage // 可为 nil
}

type inflightCall struct {
	done  chan struct{}
	token string
	err   error
}

var (
	globalTokenManager = &tokenManager{
		entries:  make(map[string]tokenEntry),
		inflight: make(map[string]*inflightCall),
	}
	tokenStorageMu sync.RWMutex
)

// SetTokenStorage 注入外部 Token 存储（如 Redis），通常在 boot 阶段调用一次
func SetTokenStorage(storage TokenStorage) {
	tokenStorageMu.Lock()
	defer tokenStorageMu.Unlock()
	globalTokenManager.storage = storage
}

// getStorage 安全读取当前 TokenStorage
func (m *tokenManager) getStorage() TokenStorage {
	tokenStorageMu.RLock()
	defer tokenStorageMu.RUnlock()
	return m.storage
}

// getFromMem 从 L1 读取有效 token
func (m *tokenManager) getFromMem(clientID string) (string, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	e, ok := m.entries[clientID]
	if !ok {
		return "", false
	}
	if time.Now().After(e.expiresAt) {
		return "", false
	}
	return e.token, true
}

// setToMem 写入 L1
func (m *tokenManager) setToMem(clientID, token string, ttl time.Duration) {
	m.mu.Lock()
	defer m.mu.Unlock()
	m.entries[clientID] = tokenEntry{
		token:     token,
		expiresAt: time.Now().Add(ttl),
	}
}

// Fetch 获取 access_token：按 L1 → L2(storage) → refresh 顺序查找，带 singleflight
// refreshFn 由调用方（dingtalk.Client）注入，返回 (token, ttl, err)
func (m *tokenManager) Fetch(
	ctx context.Context,
	clientID string,
	refreshFn func(ctx context.Context) (string, time.Duration, error),
) (string, error) {
	// L1
	if tk, ok := m.getFromMem(clientID); ok {
		return tk, nil
	}

	// L2
	if storage := m.getStorage(); storage != nil {
		if tk, err := storage.GetToken(ctx, clientID); err == nil && tk != "" {
			// 回填 L1（保守给一个最大 TTL，由 L2 过期决定真实生命周期）
			m.setToMem(clientID, tk, tokenMaxTTL-tokenRefreshBuffer)
			return tk, nil
		}
	}

	// Refresh (singleflight)
	m.inflightMu.Lock()
	call, exists := m.inflight[clientID]
	if !exists {
		call = &inflightCall{done: make(chan struct{})}
		m.inflight[clientID] = call
		m.inflightMu.Unlock()

		go func() {
			defer func() {
				m.inflightMu.Lock()
				delete(m.inflight, clientID)
				m.inflightMu.Unlock()
				close(call.done)
			}()

			token, ttl, err := refreshFn(ctx)
			if err != nil {
				call.err = err
				return
			}

			effectiveTTL := ttl - tokenRefreshBuffer
			if effectiveTTL <= 0 {
				effectiveTTL = tokenMaxTTL - tokenRefreshBuffer
			}

			m.setToMem(clientID, token, effectiveTTL)
			if storage := m.getStorage(); storage != nil {
				_ = storage.SetToken(ctx, clientID, token, effectiveTTL)
			}
			call.token = token
		}()
	} else {
		m.inflightMu.Unlock()
	}

	select {
	case <-call.done:
		return call.token, call.err
	case <-ctx.Done():
		return "", ctx.Err()
	}
}

// Invalidate 主动失效某个 clientID 的缓存（例如收到 40014 InvalidAuthentication 时）
func (m *tokenManager) Invalidate(ctx context.Context, clientID string) {
	m.mu.Lock()
	delete(m.entries, clientID)
	m.mu.Unlock()
	if storage := m.getStorage(); storage != nil {
		_ = storage.SetToken(ctx, clientID, "", time.Second)
	}
}