github.com/thanos-io/thanos@v0.32.5/internal/cortex/util/active_user.go (about)

     1  // Copyright (c) The Cortex Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package util
     5  
     6  import (
     7  	"context"
     8  	"sync"
     9  	"time"
    10  
    11  	"go.uber.org/atomic"
    12  
    13  	"github.com/thanos-io/thanos/internal/cortex/util/services"
    14  )
    15  
    16  // ActiveUsers keeps track of latest user's activity timestamp,
    17  // and allows purging users that are no longer active.
    18  type ActiveUsers struct {
    19  	mu         sync.RWMutex
    20  	timestamps map[string]*atomic.Int64 // As long as unit used by Update and Purge is the same, it doesn't matter what it is.
    21  }
    22  
    23  func NewActiveUsers() *ActiveUsers {
    24  	return &ActiveUsers{
    25  		timestamps: map[string]*atomic.Int64{},
    26  	}
    27  }
    28  
    29  func (m *ActiveUsers) UpdateUserTimestamp(userID string, ts int64) {
    30  	m.mu.RLock()
    31  	u := m.timestamps[userID]
    32  	m.mu.RUnlock()
    33  
    34  	if u != nil {
    35  		u.Store(ts)
    36  		return
    37  	}
    38  
    39  	// Pre-allocate new atomic to avoid doing allocation with lock held.
    40  	newAtomic := atomic.NewInt64(ts)
    41  
    42  	// We need RW lock to create new entry.
    43  	m.mu.Lock()
    44  	u = m.timestamps[userID]
    45  
    46  	if u != nil {
    47  		// Unlock first to reduce contention.
    48  		m.mu.Unlock()
    49  
    50  		u.Store(ts)
    51  		return
    52  	}
    53  
    54  	m.timestamps[userID] = newAtomic
    55  	m.mu.Unlock()
    56  }
    57  
    58  // PurgeInactiveUsers removes users that were last active before given deadline, and returns removed users.
    59  func (m *ActiveUsers) PurgeInactiveUsers(deadline int64) []string {
    60  	// Find inactive users with read-lock.
    61  	m.mu.RLock()
    62  	inactive := make([]string, 0, len(m.timestamps))
    63  
    64  	for userID, ts := range m.timestamps {
    65  		if ts.Load() <= deadline {
    66  			inactive = append(inactive, userID)
    67  		}
    68  	}
    69  	m.mu.RUnlock()
    70  
    71  	if len(inactive) == 0 {
    72  		return nil
    73  	}
    74  
    75  	// Cleanup inactive users.
    76  	for ix := 0; ix < len(inactive); {
    77  		userID := inactive[ix]
    78  		deleted := false
    79  
    80  		m.mu.Lock()
    81  		u := m.timestamps[userID]
    82  		if u != nil && u.Load() <= deadline {
    83  			delete(m.timestamps, userID)
    84  			deleted = true
    85  		}
    86  		m.mu.Unlock()
    87  
    88  		if deleted {
    89  			// keep it in the output
    90  			ix++
    91  		} else {
    92  			// not really inactive, remove it from output
    93  			inactive = append(inactive[:ix], inactive[ix+1:]...)
    94  		}
    95  	}
    96  
    97  	return inactive
    98  }
    99  
   100  // ActiveUsersCleanupService tracks active users, and periodically purges inactive ones while running.
   101  type ActiveUsersCleanupService struct {
   102  	services.Service
   103  
   104  	activeUsers     *ActiveUsers
   105  	cleanupFunc     func(string)
   106  	inactiveTimeout time.Duration
   107  }
   108  
   109  func NewActiveUsersCleanupWithDefaultValues(cleanupFn func(string)) *ActiveUsersCleanupService {
   110  	return NewActiveUsersCleanupService(3*time.Minute, 15*time.Minute, cleanupFn)
   111  }
   112  
   113  func NewActiveUsersCleanupService(cleanupInterval, inactiveTimeout time.Duration, cleanupFn func(string)) *ActiveUsersCleanupService {
   114  	s := &ActiveUsersCleanupService{
   115  		activeUsers:     NewActiveUsers(),
   116  		cleanupFunc:     cleanupFn,
   117  		inactiveTimeout: inactiveTimeout,
   118  	}
   119  
   120  	s.Service = services.NewTimerService(cleanupInterval, nil, s.iteration, nil).WithName("active users cleanup")
   121  	return s
   122  }
   123  
   124  func (s *ActiveUsersCleanupService) UpdateUserTimestamp(user string, now time.Time) {
   125  	s.activeUsers.UpdateUserTimestamp(user, now.UnixNano())
   126  }
   127  
   128  func (s *ActiveUsersCleanupService) iteration(_ context.Context) error {
   129  	inactiveUsers := s.activeUsers.PurgeInactiveUsers(time.Now().Add(-s.inactiveTimeout).UnixNano())
   130  	for _, userID := range inactiveUsers {
   131  		s.cleanupFunc(userID)
   132  	}
   133  	return nil
   134  }