github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/util/active_user.go (about)

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