github.com/koko1123/flow-go-1@v0.29.6/engine/execution/ingestion/uploader/manager.go (about)

     1  package uploader
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"golang.org/x/sync/errgroup"
     8  
     9  	"github.com/koko1123/flow-go-1/engine/execution"
    10  	"github.com/koko1123/flow-go-1/module"
    11  	"github.com/koko1123/flow-go-1/module/trace"
    12  )
    13  
    14  // Manager encapsulates the logic for uploading computation results to cloud storage. It
    15  type Manager struct {
    16  	enabled   bool
    17  	tracer    module.Tracer
    18  	uploaders []Uploader
    19  	mu        sync.RWMutex
    20  }
    21  
    22  // NewManager creates a new uploader manager
    23  func NewManager(tracer module.Tracer) *Manager {
    24  	return &Manager{
    25  		tracer: tracer,
    26  	}
    27  }
    28  
    29  // AddUploader adds an uploader to the manager
    30  func (m *Manager) AddUploader(uploader Uploader) {
    31  	m.mu.Lock()
    32  	defer m.mu.Unlock()
    33  
    34  	m.enabled = true
    35  	m.uploaders = append(m.uploaders, uploader)
    36  }
    37  
    38  // SetEnabled enables or disables the manager
    39  func (m *Manager) SetEnabled(enabled bool) {
    40  	m.mu.Lock()
    41  	defer m.mu.Unlock()
    42  
    43  	m.enabled = enabled
    44  }
    45  
    46  // Enabled returns whether the manager is enabled
    47  func (m *Manager) Enabled() bool {
    48  	m.mu.RLock()
    49  	defer m.mu.RUnlock()
    50  
    51  	return m.enabled
    52  }
    53  
    54  // Upload uploads the given computation result with all uploaders
    55  // Any errors returned by the uploaders may be considered benign
    56  func (m *Manager) Upload(ctx context.Context, result *execution.ComputationResult) error {
    57  	m.mu.RLock()
    58  	defer m.mu.RUnlock()
    59  
    60  	var group errgroup.Group
    61  
    62  	for _, uploader := range m.uploaders {
    63  		uploader := uploader
    64  
    65  		group.Go(func() error {
    66  			span, _ := m.tracer.StartSpanFromContext(ctx, trace.EXEUploadCollections)
    67  			defer span.End()
    68  
    69  			return uploader.Upload(result)
    70  		})
    71  	}
    72  
    73  	return group.Wait()
    74  }
    75  
    76  // RetryUploads retries uploads for all uploaders that implement RetryableUploaderWrapper
    77  // Any errors returned by the uploaders may be considered benign
    78  func (m *Manager) RetryUploads() (err error) {
    79  	m.mu.RLock()
    80  	defer m.mu.RUnlock()
    81  
    82  	for _, u := range m.uploaders {
    83  		switch retryableUploaderWraper := u.(type) {
    84  		case RetryableUploaderWrapper:
    85  			err = retryableUploaderWraper.RetryUpload()
    86  		}
    87  	}
    88  	return err
    89  }