github.com/filecoin-project/bacalhau@v0.3.23-0.20230228154132-45c989550ace/pkg/storage/parallel.go (about)

     1  package storage
     2  
     3  import (
     4  	"context"
     5  
     6  	"github.com/filecoin-project/bacalhau/pkg/model"
     7  	"github.com/filecoin-project/bacalhau/pkg/util/generic"
     8  	"go.ptx.dk/multierrgroup"
     9  )
    10  
    11  // ParallelPrepareStorage downloads all of the data necessary for the passed
    12  // storage specs in parallel, and returns a map of specs to their download
    13  // volume counterparts.
    14  func ParallelPrepareStorage(
    15  	ctx context.Context,
    16  	provider StorageProvider,
    17  	specs []model.StorageSpec,
    18  ) (map[*model.StorageSpec]StorageVolume, error) {
    19  	volumes := generic.SyncMap[*model.StorageSpec, StorageVolume]{}
    20  	waitgroup := multierrgroup.Group{}
    21  
    22  	for _, inputStorageSpec := range specs {
    23  		spec := inputStorageSpec // https://golang.org/doc/faq#closures_and_goroutines
    24  
    25  		addStorageSpec := func() error {
    26  			var storageProvider Storage
    27  			var volumeMount StorageVolume
    28  			storageProvider, err := provider.Get(ctx, spec.StorageSource)
    29  			if err != nil {
    30  				return err
    31  			}
    32  
    33  			volumeMount, err = storageProvider.PrepareStorage(ctx, spec)
    34  			if err != nil {
    35  				return err
    36  			}
    37  
    38  			volumes.Put(&spec, volumeMount)
    39  			return nil
    40  		}
    41  
    42  		waitgroup.Go(addStorageSpec)
    43  	}
    44  
    45  	err := waitgroup.Wait()
    46  
    47  	returnMap := map[*model.StorageSpec]StorageVolume{}
    48  	volumes.Iter(func(key *model.StorageSpec, value StorageVolume) bool {
    49  		returnMap[key] = value
    50  		return true
    51  	})
    52  	return returnMap, err
    53  }