github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/stores/indexshipper/uploads/table_manager.go (about)

     1  package uploads
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/go-kit/log/level"
     9  	"github.com/prometheus/client_golang/prometheus"
    10  
    11  	"github.com/grafana/loki/pkg/storage/stores/indexshipper/index"
    12  	"github.com/grafana/loki/pkg/storage/stores/indexshipper/storage"
    13  	util_log "github.com/grafana/loki/pkg/util/log"
    14  )
    15  
    16  type Config struct {
    17  	UploadInterval time.Duration
    18  	DBRetainPeriod time.Duration
    19  }
    20  
    21  type TableManager interface {
    22  	Stop()
    23  	AddIndex(tableName, userID string, index index.Index) error
    24  	ForEach(tableName, userID string, callback index.ForEachIndexCallback) error
    25  }
    26  
    27  type tableManager struct {
    28  	cfg           Config
    29  	storageClient storage.Client
    30  
    31  	tables    map[string]Table
    32  	tablesMtx sync.RWMutex
    33  	metrics   *metrics
    34  
    35  	ctx    context.Context
    36  	cancel context.CancelFunc
    37  	wg     sync.WaitGroup
    38  }
    39  
    40  func NewTableManager(cfg Config, storageClient storage.Client, reg prometheus.Registerer) (TableManager, error) {
    41  	ctx, cancel := context.WithCancel(context.Background())
    42  	tm := tableManager{
    43  		cfg:           cfg,
    44  		storageClient: storageClient,
    45  		tables:        map[string]Table{},
    46  		metrics:       newMetrics(reg),
    47  		ctx:           ctx,
    48  		cancel:        cancel,
    49  	}
    50  
    51  	go tm.loop()
    52  	return &tm, nil
    53  }
    54  
    55  func (tm *tableManager) loop() {
    56  	tm.wg.Add(1)
    57  	defer tm.wg.Done()
    58  
    59  	tm.uploadTables(context.Background())
    60  
    61  	syncTicker := time.NewTicker(tm.cfg.UploadInterval)
    62  	defer syncTicker.Stop()
    63  
    64  	for {
    65  		select {
    66  		case <-syncTicker.C:
    67  			tm.uploadTables(context.Background())
    68  		case <-tm.ctx.Done():
    69  			return
    70  		}
    71  	}
    72  }
    73  
    74  func (tm *tableManager) Stop() {
    75  	level.Info(util_log.Logger).Log("msg", "stopping table manager")
    76  
    77  	tm.cancel()
    78  	tm.wg.Wait()
    79  
    80  	tm.uploadTables(context.Background())
    81  
    82  	tm.tablesMtx.Lock()
    83  	defer tm.tablesMtx.Unlock()
    84  	for _, table := range tm.tables {
    85  		table.Stop()
    86  	}
    87  
    88  	tm.tables = map[string]Table{}
    89  }
    90  
    91  func (tm *tableManager) AddIndex(tableName, userID string, index index.Index) error {
    92  	return tm.getOrCreateTable(tableName).AddIndex(userID, index)
    93  }
    94  
    95  func (tm *tableManager) getTable(tableName string) (Table, bool) {
    96  	tm.tablesMtx.RLock()
    97  	defer tm.tablesMtx.RUnlock()
    98  
    99  	table, ok := tm.tables[tableName]
   100  	return table, ok
   101  }
   102  
   103  func (tm *tableManager) getOrCreateTable(tableName string) Table {
   104  	table, ok := tm.getTable(tableName)
   105  	if ok {
   106  		return table
   107  	}
   108  
   109  	tm.tablesMtx.Lock()
   110  	defer tm.tablesMtx.Unlock()
   111  
   112  	table, ok = tm.tables[tableName]
   113  	if !ok {
   114  		table = NewTable(tableName, tm.storageClient)
   115  		tm.tables[tableName] = table
   116  	}
   117  
   118  	return table
   119  }
   120  
   121  func (tm *tableManager) ForEach(tableName, userID string, callback index.ForEachIndexCallback) error {
   122  	table, ok := tm.getTable(tableName)
   123  	if !ok {
   124  		return nil
   125  	}
   126  
   127  	return table.ForEach(userID, callback)
   128  }
   129  
   130  func (tm *tableManager) uploadTables(ctx context.Context) {
   131  	tm.tablesMtx.RLock()
   132  	defer tm.tablesMtx.RUnlock()
   133  
   134  	level.Info(util_log.Logger).Log("msg", "uploading tables")
   135  
   136  	status := statusSuccess
   137  	for _, table := range tm.tables {
   138  		err := table.Upload(ctx)
   139  		if err != nil {
   140  			status = statusFailure
   141  			level.Error(util_log.Logger).Log("msg", "failed to upload table", "table", table.Name(), "err", err)
   142  			continue
   143  		}
   144  
   145  		// cleanup uploaded dbs from local disk after retain period
   146  		err = table.Cleanup(tm.cfg.DBRetainPeriod)
   147  		if err != nil {
   148  			// we do not want to stop uploading of dbs due to failures in cleaning them up so logging just the error here.
   149  			level.Error(util_log.Logger).Log("msg", "failed to cleanup uploaded index past their retention period", "table", table.Name(), "err", err)
   150  		}
   151  	}
   152  
   153  	tm.metrics.tablesUploadOperationTotal.WithLabelValues(status).Inc()
   154  }