github.com/Ilhicas/nomad@v1.0.4-0.20210304152020-e86851182bc3/client/allocrunner/taskrunner/artifact_hook.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  
     7  	log "github.com/hashicorp/go-hclog"
     8  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
     9  	"github.com/hashicorp/nomad/client/allocrunner/taskrunner/getter"
    10  	ti "github.com/hashicorp/nomad/client/allocrunner/taskrunner/interfaces"
    11  	"github.com/hashicorp/nomad/nomad/structs"
    12  )
    13  
    14  // artifactHook downloads artifacts for a task.
    15  type artifactHook struct {
    16  	eventEmitter ti.EventEmitter
    17  	logger       log.Logger
    18  }
    19  
    20  func newArtifactHook(e ti.EventEmitter, logger log.Logger) *artifactHook {
    21  	h := &artifactHook{
    22  		eventEmitter: e,
    23  	}
    24  	h.logger = logger.Named(h.Name())
    25  	return h
    26  }
    27  
    28  func (*artifactHook) Name() string {
    29  	// Copied in client/state when upgrading from <0.9 schemas, so if you
    30  	// change it here you also must change it there.
    31  	return "artifacts"
    32  }
    33  
    34  func (h *artifactHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
    35  	if len(req.Task.Artifacts) == 0 {
    36  		resp.Done = true
    37  		return nil
    38  	}
    39  
    40  	// Initialize hook state to store download progress
    41  	resp.State = make(map[string]string, len(req.Task.Artifacts))
    42  
    43  	h.eventEmitter.EmitEvent(structs.NewTaskEvent(structs.TaskDownloadingArtifacts))
    44  
    45  	for _, artifact := range req.Task.Artifacts {
    46  		aid := artifact.Hash()
    47  		if req.PreviousState[aid] != "" {
    48  			h.logger.Trace("skipping already downloaded artifact", "artifact", artifact.GetterSource)
    49  			resp.State[aid] = req.PreviousState[aid]
    50  			continue
    51  		}
    52  
    53  		h.logger.Debug("downloading artifact", "artifact", artifact.GetterSource)
    54  		//XXX add ctx to GetArtifact to allow cancelling long downloads
    55  		if err := getter.GetArtifact(req.TaskEnv, artifact); err != nil {
    56  
    57  			wrapped := structs.NewRecoverableError(
    58  				fmt.Errorf("failed to download artifact %q: %v", artifact.GetterSource, err),
    59  				true,
    60  			)
    61  			herr := NewHookError(wrapped, structs.NewTaskEvent(structs.TaskArtifactDownloadFailed).SetDownloadError(wrapped))
    62  
    63  			return herr
    64  		}
    65  
    66  		// Mark artifact as downloaded to avoid re-downloading due to
    67  		// retries caused by subsequent artifacts failing. Any
    68  		// non-empty value works.
    69  		resp.State[aid] = "1"
    70  	}
    71  
    72  	resp.Done = true
    73  	return nil
    74  }