geo/server/internal/tenant/app/desktop_presence.go

package app

import (
	"context"
	"strconv"
	"time"

	"github.com/google/uuid"
	goredis "github.com/redis/go-redis/v9"
)

// desktopClientPresenceTTL is also used as the monitoring primary-client lease TTL.
// Desktop heartbeat interval is TTL/2, so unexpected client exits are reflected
// as offline within one TTL even when the explicit offline request is missed.
const desktopClientPresenceTTL = 30 * time.Second
const desktopAccountPresenceTTL = 30 * time.Second
const desktopTaskConsumerPresenceTTL = 75 * time.Second

func desktopClientPresenceKey(clientID uuid.UUID) string {
	return "desktop:presence:client:" + clientID.String()
}

func desktopTaskConsumerPresenceKey(clientID uuid.UUID) string {
	return "desktop:presence:task_consumer:" + clientID.String()
}

func desktopAccountPresenceKey(accountID uuid.UUID) string {
	return "desktop:presence:account:" + accountID.String()
}

func desktopClientAccountsPresenceKey(clientID uuid.UUID) string {
	return "desktop:presence:client_accounts:" + clientID.String()
}

func desktopUserClientsPresenceKey(tenantID, workspaceID, userID int64) string {
	return "desktop:presence:user_clients:" +
		strconv.FormatInt(tenantID, 10) + ":" +
		strconv.FormatInt(workspaceID, 10) + ":" +
		strconv.FormatInt(userID, 10)
}

func markDesktopClientPresent(ctx context.Context, redis *goredis.Client, clientID uuid.UUID) {
	if redis == nil {
		return
	}

	_ = redis.Set(ctx, desktopClientPresenceKey(clientID), time.Now().UTC().Format(time.RFC3339), desktopClientPresenceTTL).Err()
}

func markDesktopTaskConsumerPresent(ctx context.Context, redis *goredis.Client, clientID uuid.UUID) {
	if redis == nil || clientID == uuid.Nil {
		return
	}

	_ = redis.Set(ctx, desktopTaskConsumerPresenceKey(clientID), time.Now().UTC().Format(time.RFC3339), desktopTaskConsumerPresenceTTL).Err()
}

func MarkDesktopTaskConsumerPresent(ctx context.Context, redis *goredis.Client, clientID uuid.UUID) {
	markDesktopTaskConsumerPresent(ctx, redis, clientID)
}

func markDesktopClientPresentForUser(ctx context.Context, redis *goredis.Client, clientID uuid.UUID, tenantID, workspaceID, userID int64) {
	markDesktopClientPresent(ctx, redis, clientID)
	if redis == nil || clientID == uuid.Nil || tenantID == 0 || workspaceID == 0 || userID == 0 {
		return
	}

	now := time.Now().UTC()
	staleBefore := strconv.FormatInt(now.Add(-desktopClientPresenceTTL).UnixMilli(), 10)
	key := desktopUserClientsPresenceKey(tenantID, workspaceID, userID)
	pipe := redis.Pipeline()
	pipe.ZRemRangeByScore(ctx, key, "-inf", staleBefore)
	pipe.ZAdd(ctx, key, goredis.Z{
		Score:  float64(now.UnixMilli()),
		Member: clientID.String(),
	})
	pipe.Expire(ctx, key, desktopClientPresenceTTL*4)
	_, _ = pipe.Exec(ctx)
}

func clearDesktopClientPresence(ctx context.Context, redis *goredis.Client, clientID uuid.UUID) {
	if redis == nil {
		return
	}

	_ = redis.Del(ctx, desktopClientPresenceKey(clientID), desktopTaskConsumerPresenceKey(clientID)).Err()
}

func clearDesktopClientPresenceStateForUser(ctx context.Context, redis *goredis.Client, clientID uuid.UUID, tenantID, workspaceID, userID int64) {
	clearDesktopClientPresenceState(ctx, redis, clientID)
	if redis == nil || clientID == uuid.Nil || tenantID == 0 || workspaceID == 0 || userID == 0 {
		return
	}

	staleBefore := strconv.FormatInt(time.Now().UTC().Add(-desktopClientPresenceTTL).UnixMilli(), 10)
	key := desktopUserClientsPresenceKey(tenantID, workspaceID, userID)
	pipe := redis.Pipeline()
	pipe.ZRem(ctx, key, clientID.String())
	pipe.ZRemRangeByScore(ctx, key, "-inf", staleBefore)
	_, _ = pipe.Exec(ctx)
}

func loadDesktopClientPresence(ctx context.Context, redis *goredis.Client, clientIDs []uuid.UUID) map[uuid.UUID]bool {
	if redis == nil || len(clientIDs) == 0 {
		return nil
	}

	uniqueIDs := make([]uuid.UUID, 0, len(clientIDs))
	seen := make(map[uuid.UUID]struct{}, len(clientIDs))
	for _, clientID := range clientIDs {
		if _, ok := seen[clientID]; ok {
			continue
		}
		seen[clientID] = struct{}{}
		uniqueIDs = append(uniqueIDs, clientID)
	}

	pipe := redis.Pipeline()
	cmds := make(map[uuid.UUID]*goredis.IntCmd, len(uniqueIDs))
	for _, clientID := range uniqueIDs {
		cmds[clientID] = pipe.Exists(ctx, desktopClientPresenceKey(clientID))
	}

	if _, err := pipe.Exec(ctx); err != nil && err != goredis.Nil {
		return nil
	}

	result := make(map[uuid.UUID]bool, len(cmds))
	for clientID, cmd := range cmds {
		result[clientID] = cmd.Val() > 0
	}
	return result
}

func loadDesktopTaskConsumerPresence(ctx context.Context, redis *goredis.Client, clientIDs []uuid.UUID) map[uuid.UUID]bool {
	if redis == nil || len(clientIDs) == 0 {
		return nil
	}

	uniqueIDs := make([]uuid.UUID, 0, len(clientIDs))
	seen := make(map[uuid.UUID]struct{}, len(clientIDs))
	for _, clientID := range clientIDs {
		if clientID == uuid.Nil {
			continue
		}
		if _, ok := seen[clientID]; ok {
			continue
		}
		seen[clientID] = struct{}{}
		uniqueIDs = append(uniqueIDs, clientID)
	}
	if len(uniqueIDs) == 0 {
		return nil
	}

	pipe := redis.Pipeline()
	cmds := make(map[uuid.UUID]*goredis.IntCmd, len(uniqueIDs))
	for _, clientID := range uniqueIDs {
		cmds[clientID] = pipe.Exists(ctx, desktopTaskConsumerPresenceKey(clientID))
	}

	if _, err := pipe.Exec(ctx); err != nil && err != goredis.Nil {
		return nil
	}

	result := make(map[uuid.UUID]bool, len(cmds))
	for clientID, cmd := range cmds {
		result[clientID] = cmd.Val() > 0
	}
	return result
}

func loadDesktopUserOnlineClientCount(ctx context.Context, redis *goredis.Client, tenantID, workspaceID, userID int64, now time.Time) (int, bool) {
	if redis == nil || tenantID == 0 || workspaceID == 0 || userID == 0 {
		return 0, false
	}

	staleBefore := strconv.FormatInt(now.UTC().Add(-desktopClientPresenceTTL).UnixMilli(), 10)
	key := desktopUserClientsPresenceKey(tenantID, workspaceID, userID)
	pipe := redis.Pipeline()
	pipe.ZRemRangeByScore(ctx, key, "-inf", staleBefore)
	countCmd := pipe.ZCard(ctx, key)
	if _, err := pipe.Exec(ctx); err != nil && err != goredis.Nil {
		return 0, false
	}
	return int(countCmd.Val()), true
}

func markDesktopAccountsPresent(ctx context.Context, redis *goredis.Client, clientID uuid.UUID, accountIDs []uuid.UUID) {
	if redis == nil {
		return
	}

	uniqueIDs := uniqueUUIDs(accountIDs)
	trackedAccountIDs := loadTrackedDesktopClientAccountIDs(ctx, redis, clientID)
	nextAccountSet := make(map[uuid.UUID]struct{}, len(uniqueIDs))
	for _, accountID := range uniqueIDs {
		nextAccountSet[accountID] = struct{}{}
	}

	now := float64(time.Now().UTC().UnixMilli())
	staleBefore := strconv.FormatInt(time.Now().UTC().Add(-desktopAccountPresenceTTL).UnixMilli(), 10)
	clientAccountsKey := desktopClientAccountsPresenceKey(clientID)
	clientMember := clientID.String()

	pipe := redis.Pipeline()
	pipe.Del(ctx, clientAccountsKey)
	if len(uniqueIDs) > 0 {
		members := make([]interface{}, 0, len(uniqueIDs))
		for _, accountID := range uniqueIDs {
			members = append(members, accountID.String())
		}
		pipe.SAdd(ctx, clientAccountsKey, members...)
		pipe.Expire(ctx, clientAccountsKey, desktopAccountPresenceTTL*4)
	}

	for _, accountID := range trackedAccountIDs {
		if _, ok := nextAccountSet[accountID]; ok {
			continue
		}
		pipe.ZRem(ctx, desktopAccountPresenceKey(accountID), clientMember)
	}
	for _, accountID := range uniqueIDs {
		key := desktopAccountPresenceKey(accountID)
		pipe.ZRemRangeByScore(ctx, key, "-inf", staleBefore)
		pipe.ZAdd(ctx, key, goredis.Z{
			Score:  now,
			Member: clientMember,
		})
		pipe.Expire(ctx, key, desktopAccountPresenceTTL*4)
	}
	_, _ = pipe.Exec(ctx)

	cleanupDesktopAccountPresenceKeys(ctx, redis, trackedAccountIDs)
}

func clearDesktopClientPresenceState(ctx context.Context, redis *goredis.Client, clientID uuid.UUID) {
	if redis == nil {
		return
	}

	clearDesktopClientPresence(ctx, redis, clientID)

	trackedAccountIDs := loadTrackedDesktopClientAccountIDs(ctx, redis, clientID)
	if len(trackedAccountIDs) == 0 {
		_ = redis.Del(ctx, desktopClientAccountsPresenceKey(clientID)).Err()
		return
	}

	clientMember := clientID.String()
	pipe := redis.Pipeline()
	for _, accountID := range trackedAccountIDs {
		pipe.ZRem(ctx, desktopAccountPresenceKey(accountID), clientMember)
	}
	pipe.Del(ctx, desktopClientAccountsPresenceKey(clientID))
	_, _ = pipe.Exec(ctx)

	cleanupDesktopAccountPresenceKeys(ctx, redis, trackedAccountIDs)
}

func loadDesktopAccountPresence(ctx context.Context, redis *goredis.Client, accountIDs []uuid.UUID) map[uuid.UUID]uuid.UUID {
	if redis == nil || len(accountIDs) == 0 {
		return nil
	}

	uniqueIDs := uniqueUUIDs(accountIDs)
	staleBefore := strconv.FormatInt(time.Now().UTC().Add(-desktopAccountPresenceTTL).UnixMilli(), 10)

	pipe := redis.Pipeline()
	commands := make(map[uuid.UUID]*goredis.ZSliceCmd, len(uniqueIDs))
	for _, accountID := range uniqueIDs {
		key := desktopAccountPresenceKey(accountID)
		pipe.ZRemRangeByScore(ctx, key, "-inf", staleBefore)
		commands[accountID] = pipe.ZRevRangeWithScores(ctx, key, 0, 0)
	}

	if _, err := pipe.Exec(ctx); err != nil && err != goredis.Nil {
		return nil
	}

	result := make(map[uuid.UUID]uuid.UUID, len(commands))
	for accountID, cmd := range commands {
		values := cmd.Val()
		if len(values) == 0 {
			continue
		}
		clientIDText, ok := values[0].Member.(string)
		if !ok || clientIDText == "" {
			continue
		}
		clientID, err := uuid.Parse(clientIDText)
		if err != nil {
			continue
		}
		result[accountID] = clientID
	}
	return result
}

func uniqueUUIDs(values []uuid.UUID) []uuid.UUID {
	if len(values) == 0 {
		return nil
	}

	unique := make([]uuid.UUID, 0, len(values))
	seen := make(map[uuid.UUID]struct{}, len(values))
	for _, value := range values {
		if _, ok := seen[value]; ok {
			continue
		}
		seen[value] = struct{}{}
		unique = append(unique, value)
	}
	return unique
}

func loadTrackedDesktopClientAccountIDs(ctx context.Context, redis *goredis.Client, clientID uuid.UUID) []uuid.UUID {
	if redis == nil {
		return nil
	}

	values, err := redis.SMembers(ctx, desktopClientAccountsPresenceKey(clientID)).Result()
	if err != nil || len(values) == 0 {
		return nil
	}

	accountIDs := make([]uuid.UUID, 0, len(values))
	for _, value := range values {
		accountID, parseErr := uuid.Parse(value)
		if parseErr != nil {
			continue
		}
		accountIDs = append(accountIDs, accountID)
	}
	return uniqueUUIDs(accountIDs)
}

func cleanupDesktopAccountPresenceKeys(ctx context.Context, redis *goredis.Client, accountIDs []uuid.UUID) {
	if redis == nil || len(accountIDs) == 0 {
		return
	}

	uniqueIDs := uniqueUUIDs(accountIDs)
	pipe := redis.Pipeline()
	commands := make(map[uuid.UUID]*goredis.IntCmd, len(uniqueIDs))
	for _, accountID := range uniqueIDs {
		commands[accountID] = pipe.ZCard(ctx, desktopAccountPresenceKey(accountID))
	}
	if _, err := pipe.Exec(ctx); err != nil && err != goredis.Nil {
		return
	}

	deletePipe := redis.Pipeline()
	shouldDelete := false
	for accountID, cmd := range commands {
		if cmd.Val() > 0 {
			continue
		}
		deletePipe.Del(ctx, desktopAccountPresenceKey(accountID))
		shouldDelete = true
	}
	if shouldDelete {
		_, _ = deletePipe.Exec(ctx)
	}
}