github.com/grafana/pyroscope@v1.18.0/pkg/phlaredb/bucket/tenant_scanner.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/cortexproject/cortex/blob/master/pkg/storage/tsdb/users_scanner.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The Cortex Authors.
     5  
     6  package bucket
     7  
     8  import (
     9  	"context"
    10  
    11  	"github.com/go-kit/log"
    12  	"github.com/go-kit/log/level"
    13  
    14  	"github.com/grafana/pyroscope/pkg/objstore"
    15  )
    16  
    17  // AllTenants returns true to each call and should be used whenever the UsersScanner should not filter out
    18  // any user due to sharding.
    19  func AllTenants(_ string) (bool, error) {
    20  	return true, nil
    21  }
    22  
    23  type TenantsScanner struct {
    24  	bucketClient objstore.Bucket
    25  	logger       log.Logger
    26  	isOwned      func(userID string) (bool, error)
    27  }
    28  
    29  func NewTenantsScanner(bucketClient objstore.Bucket, isOwned func(userID string) (bool, error), logger log.Logger) *TenantsScanner {
    30  	return &TenantsScanner{
    31  		bucketClient: bucketClient,
    32  		logger:       logger,
    33  		isOwned:      isOwned,
    34  	}
    35  }
    36  
    37  // ScanTenants returns a fresh list of users found in the storage, that are not marked for deletion,
    38  // and list of users marked for deletion.
    39  //
    40  // If sharding is enabled, returned lists contains only the users owned by this instance.
    41  func (s *TenantsScanner) ScanTenants(ctx context.Context) (users, markedForDeletion []string, err error) {
    42  	users, err = ListUsers(ctx, s.bucketClient)
    43  	if err != nil {
    44  		return nil, nil, err
    45  	}
    46  
    47  	// Check users for being owned by instance, and split users into non-deleted and deleted.
    48  	// We do these checks after listing all users, to improve cacheability of Iter (result is only cached at the end of Iter call).
    49  	for ix := 0; ix < len(users); {
    50  		userID := users[ix]
    51  
    52  		// Check if it's owned by this instance.
    53  		owned, err := s.isOwned(userID)
    54  		if err != nil {
    55  			level.Warn(s.logger).Log("msg", "unable to check if user is owned by this shard", "tenant", userID, "err", err)
    56  		} else if !owned {
    57  			users = append(users[:ix], users[ix+1:]...)
    58  			continue
    59  		}
    60  
    61  		deletionMarkExists, err := TenantDeletionMarkExists(ctx, s.bucketClient, userID)
    62  		if err != nil {
    63  			level.Warn(s.logger).Log("msg", "unable to check if user is marked for deletion", "tenant", userID, "err", err)
    64  		} else if deletionMarkExists {
    65  			users = append(users[:ix], users[ix+1:]...)
    66  			markedForDeletion = append(markedForDeletion, userID)
    67  			continue
    68  		}
    69  
    70  		ix++
    71  	}
    72  
    73  	return users, markedForDeletion, nil
    74  }