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