github.com/pf-qiu/concourse/v6@v6.7.3-0.20201207032516-1f455d73275f/atc/worker/artifact_source.go (about)

     1  package worker
     2  
     3  import (
     4  	"archive/tar"
     5  	"code.cloudfoundry.org/lager/lagerctx"
     6  	"context"
     7  	"io"
     8  	"time"
     9  
    10  	"code.cloudfoundry.org/lager"
    11  	"github.com/pf-qiu/concourse/v6/atc/compression"
    12  	"github.com/pf-qiu/concourse/v6/atc/metric"
    13  	"github.com/pf-qiu/concourse/v6/atc/runtime"
    14  	"github.com/pf-qiu/concourse/v6/tracing"
    15  	"github.com/hashicorp/go-multierror"
    16  )
    17  
    18  //go:generate counterfeiter . ArtifactSource
    19  
    20  type ArtifactSource interface {
    21  	// ExistsOn attempts to locate a volume equivalent to this source on the
    22  	// given worker. If a volume can be found, it will be used directly. If not,
    23  	// `StreamTo` will be used to copy the data to the destination instead.
    24  	ExistsOn(lager.Logger, Worker) (Volume, bool, error)
    25  }
    26  
    27  //go:generate counterfeiter . StreamableArtifactSource
    28  
    29  // Source represents data produced by the steps, that can be transferred to
    30  // other steps.
    31  type StreamableArtifactSource interface {
    32  	ArtifactSource
    33  	// StreamTo copies the data from the source to the destination. Note that
    34  	// this potentially uses a lot of network transfer, for larger artifacts, as
    35  	// the ATC will effectively act as a middleman.
    36  	StreamTo(context.Context, ArtifactDestination) error
    37  
    38  	// StreamFile returns the contents of a single file in the artifact source.
    39  	// This is used for loading a task's configuration at runtime.
    40  	StreamFile(context.Context, string) (io.ReadCloser, error)
    41  }
    42  
    43  type artifactSource struct {
    44  	artifact            runtime.Artifact
    45  	volume              Volume
    46  	compression         compression.Compression
    47  	enabledP2pStreaming bool
    48  	p2pStreamingTimeout time.Duration
    49  }
    50  
    51  func NewStreamableArtifactSource(
    52  	artifact runtime.Artifact,
    53  	volume Volume,
    54  	compression compression.Compression,
    55  	enabledP2pStreaming bool,
    56  	p2pStreamingTimeout time.Duration,
    57  ) StreamableArtifactSource {
    58  	return &artifactSource{
    59  		artifact:            artifact,
    60  		volume:              volume,
    61  		compression:         compression,
    62  		enabledP2pStreaming: enabledP2pStreaming,
    63  		p2pStreamingTimeout: p2pStreamingTimeout,
    64  	}
    65  }
    66  
    67  func (source *artifactSource) StreamTo(
    68  	ctx context.Context,
    69  	destination ArtifactDestination,
    70  ) error {
    71  	logger := lagerctx.FromContext(ctx).Session("stream-to")
    72  	logger.Info("start")
    73  	defer logger.Info("end")
    74  
    75  	ctx, span := tracing.StartSpan(ctx, "artifactSource.StreamTo", nil)
    76  	defer span.End()
    77  
    78  	var err error
    79  	if !source.enabledP2pStreaming {
    80  		err = source.streamTo(ctx, destination)
    81  	} else {
    82  		err = source.p2pStreamTo(ctx, destination)
    83  	}
    84  
    85  	// Inc counter if no error occurred.
    86  	if err == nil {
    87  		metric.Metrics.VolumesStreamed.Inc()
    88  	}
    89  
    90  	return err
    91  }
    92  
    93  func (source *artifactSource) streamTo(
    94  	ctx context.Context,
    95  	destination ArtifactDestination,
    96  ) error {
    97  	_, outSpan := tracing.StartSpan(ctx, "volume.StreamOut", tracing.Attrs{
    98  		"origin-volume": source.volume.Handle(),
    99  		"origin-worker": source.volume.WorkerName(),
   100  	})
   101  	defer outSpan.End()
   102  	out, err := source.volume.StreamOut(ctx, ".", source.compression.Encoding())
   103  
   104  	if err != nil {
   105  		tracing.End(outSpan, err)
   106  		return err
   107  	}
   108  
   109  	defer out.Close()
   110  
   111  	return destination.StreamIn(ctx, ".", source.compression.Encoding(), out)
   112  }
   113  
   114  func (source *artifactSource) p2pStreamTo(
   115  	ctx context.Context,
   116  	destination ArtifactDestination,
   117  ) error {
   118  	getCtx, getCancel := context.WithTimeout(ctx, 5*time.Second)
   119  	defer getCancel()
   120  	streamInUrl, err := destination.GetStreamInP2pUrl(getCtx, ".")
   121  	if err != nil {
   122  		return err
   123  	}
   124  
   125  	select {
   126  	case <-ctx.Done():
   127  		return context.Canceled
   128  	default:
   129  	}
   130  
   131  	_, outSpan := tracing.StartSpan(ctx, "volume.P2pStreamOut", tracing.Attrs{
   132  		"origin-volume": source.volume.Handle(),
   133  		"origin-worker": source.volume.WorkerName(),
   134  		"stream-in-url": streamInUrl,
   135  	})
   136  	defer outSpan.End()
   137  
   138  	putCtx, putCancel := context.WithTimeout(ctx, source.p2pStreamingTimeout)
   139  	defer putCancel()
   140  	return source.volume.StreamP2pOut(putCtx, ".", streamInUrl, source.compression.Encoding())
   141  }
   142  
   143  func (source *artifactSource) StreamFile(
   144  	ctx context.Context,
   145  	filepath string,
   146  ) (io.ReadCloser, error) {
   147  	out, err := source.volume.StreamOut(ctx, filepath, source.compression.Encoding())
   148  	if err != nil {
   149  		return nil, err
   150  	}
   151  
   152  	compressionReader, err := source.compression.NewReader(out)
   153  	if err != nil {
   154  		return nil, err
   155  	}
   156  	tarReader := tar.NewReader(compressionReader)
   157  
   158  	_, err = tarReader.Next()
   159  	if err != nil {
   160  		return nil, err
   161  	}
   162  
   163  	return fileReadMultiCloser{
   164  		reader: tarReader,
   165  		closers: []io.Closer{
   166  			out,
   167  			compressionReader,
   168  		},
   169  	}, nil
   170  }
   171  
   172  // Returns volume if it belongs to the worker
   173  //  otherwise, if the volume has a Resource Cache
   174  //  it checks the worker for a local volume corresponding to the Resource Cache.
   175  //  Note: The returned volume may have a different handle than the ArtifactSource's inner volume handle.
   176  func (source *artifactSource) ExistsOn(logger lager.Logger, worker Worker) (Volume, bool, error) {
   177  	if source.volume.WorkerName() == worker.Name() {
   178  		return source.volume, true, nil
   179  	}
   180  
   181  	resourceCache, found, err := worker.FindResourceCacheForVolume(source.volume)
   182  	if err != nil {
   183  		return nil, false, err
   184  	}
   185  	if found {
   186  		return worker.FindVolumeForResourceCache(logger, resourceCache)
   187  	} else {
   188  		return nil, false, nil
   189  	}
   190  
   191  }
   192  
   193  type cacheArtifactSource struct {
   194  	runtime.CacheArtifact
   195  }
   196  
   197  func NewCacheArtifactSource(artifact runtime.CacheArtifact) ArtifactSource {
   198  	return &cacheArtifactSource{artifact}
   199  }
   200  
   201  func (source *cacheArtifactSource) ExistsOn(logger lager.Logger, worker Worker) (Volume, bool, error) {
   202  	return worker.FindVolumeForTaskCache(logger, source.TeamID, source.JobID, source.StepName, source.Path)
   203  }
   204  
   205  type fileReadMultiCloser struct {
   206  	reader  io.Reader
   207  	closers []io.Closer
   208  }
   209  
   210  func (frc fileReadMultiCloser) Read(p []byte) (n int, err error) {
   211  	return frc.reader.Read(p)
   212  }
   213  
   214  func (frc fileReadMultiCloser) Close() error {
   215  	var closeErrors error
   216  
   217  	for _, closer := range frc.closers {
   218  		err := closer.Close()
   219  		if err != nil {
   220  			closeErrors = multierror.Append(closeErrors, err)
   221  		}
   222  	}
   223  
   224  	return closeErrors
   225  }