github.com/grafana/pyroscope@v1.18.0/pkg/util/active_user.go (about)

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