480 lines
14 KiB
Go
480 lines
14 KiB
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
|
|
|
|
type desktopAccountPresenceClient struct {
|
|
ClientID uuid.UUID
|
|
OnlineSinceUnixMilli int64
|
|
}
|
|
|
|
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
|
|
}
|
|
|
|
key := desktopClientPresenceKey(clientID)
|
|
now := strconv.FormatInt(time.Now().UTC().UnixMilli(), 10)
|
|
set, err := redis.SetNX(ctx, key, now, desktopClientPresenceTTL).Result()
|
|
if err != nil || set {
|
|
return
|
|
}
|
|
_ = redis.Expire(ctx, key, 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 := time.Now().UTC()
|
|
onlineSinceUnixMilli := now.UnixMilli()
|
|
if value, err := redis.Get(ctx, desktopClientPresenceKey(clientID)).Result(); err == nil {
|
|
onlineSinceUnixMilli = parseDesktopPresenceUnixMilli(value, now)
|
|
}
|
|
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.ZAdd(ctx, key, goredis.Z{
|
|
Score: float64(onlineSinceUnixMilli),
|
|
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 {
|
|
clientsByAccount := loadDesktopAccountPresenceClients(ctx, redis, accountIDs)
|
|
if len(clientsByAccount) == 0 {
|
|
return nil
|
|
}
|
|
|
|
result := make(map[uuid.UUID]uuid.UUID, len(clientsByAccount))
|
|
for accountID, clients := range clientsByAccount {
|
|
if len(clients) == 0 {
|
|
continue
|
|
}
|
|
result[accountID] = clients[len(clients)-1].ClientID
|
|
}
|
|
return result
|
|
}
|
|
|
|
func loadDesktopAccountPresenceClients(ctx context.Context, redis *goredis.Client, accountIDs []uuid.UUID) map[uuid.UUID][]desktopAccountPresenceClient {
|
|
if redis == nil || len(accountIDs) == 0 {
|
|
return nil
|
|
}
|
|
|
|
uniqueIDs := uniqueUUIDs(accountIDs)
|
|
|
|
pipe := redis.Pipeline()
|
|
commands := make(map[uuid.UUID]*goredis.ZSliceCmd, len(uniqueIDs))
|
|
for _, accountID := range uniqueIDs {
|
|
key := desktopAccountPresenceKey(accountID)
|
|
commands[accountID] = pipe.ZRangeWithScores(ctx, key, 0, -1)
|
|
}
|
|
|
|
if _, err := pipe.Exec(ctx); err != nil && err != goredis.Nil {
|
|
return nil
|
|
}
|
|
|
|
result := make(map[uuid.UUID][]desktopAccountPresenceClient, len(commands))
|
|
clientIDs := make([]uuid.UUID, 0)
|
|
staleMembersByAccount := make(map[uuid.UUID][]interface{})
|
|
for accountID, cmd := range commands {
|
|
values := cmd.Val()
|
|
if len(values) == 0 {
|
|
continue
|
|
}
|
|
seenClients := make(map[uuid.UUID]struct{}, len(values))
|
|
for _, value := range values {
|
|
clientIDText, ok := value.Member.(string)
|
|
if !ok || clientIDText == "" {
|
|
staleMembersByAccount[accountID] = append(staleMembersByAccount[accountID], value.Member)
|
|
continue
|
|
}
|
|
clientID, err := uuid.Parse(clientIDText)
|
|
if err != nil {
|
|
staleMembersByAccount[accountID] = append(staleMembersByAccount[accountID], clientIDText)
|
|
continue
|
|
}
|
|
if _, exists := seenClients[clientID]; exists {
|
|
continue
|
|
}
|
|
seenClients[clientID] = struct{}{}
|
|
result[accountID] = append(result[accountID], desktopAccountPresenceClient{
|
|
ClientID: clientID,
|
|
OnlineSinceUnixMilli: int64(value.Score),
|
|
})
|
|
clientIDs = append(clientIDs, clientID)
|
|
}
|
|
}
|
|
|
|
clientIDs = uniqueUUIDs(clientIDs)
|
|
if len(clientIDs) == 0 {
|
|
pruneDesktopAccountPresenceMembers(ctx, redis, staleMembersByAccount)
|
|
cleanupDesktopAccountPresenceKeys(ctx, redis, uniqueIDs)
|
|
return nil
|
|
}
|
|
|
|
presencePipe := redis.Pipeline()
|
|
presenceCommands := make(map[uuid.UUID]*goredis.IntCmd, len(clientIDs))
|
|
for _, clientID := range clientIDs {
|
|
presenceCommands[clientID] = presencePipe.Exists(ctx, desktopClientPresenceKey(clientID))
|
|
}
|
|
if _, err := presencePipe.Exec(ctx); err != nil && err != goredis.Nil {
|
|
return nil
|
|
}
|
|
|
|
for accountID, clients := range result {
|
|
kept := clients[:0]
|
|
for _, client := range clients {
|
|
cmd := presenceCommands[client.ClientID]
|
|
if cmd != nil && cmd.Val() > 0 {
|
|
kept = append(kept, client)
|
|
continue
|
|
}
|
|
staleMembersByAccount[accountID] = append(staleMembersByAccount[accountID], client.ClientID.String())
|
|
}
|
|
if len(kept) == 0 {
|
|
delete(result, accountID)
|
|
continue
|
|
}
|
|
result[accountID] = kept
|
|
}
|
|
|
|
pruneDesktopAccountPresenceMembers(ctx, redis, staleMembersByAccount)
|
|
cleanupDesktopAccountPresenceKeys(ctx, redis, uniqueIDs)
|
|
return result
|
|
}
|
|
|
|
func parseDesktopPresenceUnixMilli(value string, fallback time.Time) int64 {
|
|
return parseDesktopPresenceTime(value, fallback).UnixMilli()
|
|
}
|
|
|
|
func parseDesktopPresenceTime(value string, fallback time.Time) time.Time {
|
|
if millis, err := strconv.ParseInt(value, 10, 64); err == nil && millis > 0 {
|
|
return time.UnixMilli(millis).UTC()
|
|
}
|
|
if parsed, err := time.Parse(time.RFC3339Nano, value); err == nil {
|
|
return parsed.UTC()
|
|
}
|
|
if parsed, err := time.Parse(time.RFC3339, value); err == nil {
|
|
return parsed.UTC()
|
|
}
|
|
return fallback.UTC()
|
|
}
|
|
|
|
func pruneDesktopAccountPresenceMembers(ctx context.Context, redis *goredis.Client, membersByAccount map[uuid.UUID][]interface{}) {
|
|
if redis == nil || len(membersByAccount) == 0 {
|
|
return
|
|
}
|
|
|
|
pipe := redis.Pipeline()
|
|
shouldExec := false
|
|
for accountID, members := range membersByAccount {
|
|
if len(members) == 0 {
|
|
continue
|
|
}
|
|
pipe.ZRem(ctx, desktopAccountPresenceKey(accountID), members...)
|
|
shouldExec = true
|
|
}
|
|
if shouldExec {
|
|
_, _ = pipe.Exec(ctx)
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|
|
}
|