github.com/iqoqo/nomad@v0.11.3-0.20200911112621-d7021c74d101/client/allocrunner/taskrunner/dispatch_hook.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"io/ioutil"
     6  	"os"
     7  	"path/filepath"
     8  
     9  	"github.com/golang/snappy"
    10  	hclog "github.com/hashicorp/go-hclog"
    11  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
    12  	"github.com/hashicorp/nomad/nomad/structs"
    13  )
    14  
    15  // dispatchHook writes a dispatch payload to the task dir
    16  type dispatchHook struct {
    17  	payload []byte
    18  
    19  	logger hclog.Logger
    20  }
    21  
    22  func newDispatchHook(alloc *structs.Allocation, logger hclog.Logger) *dispatchHook {
    23  	h := &dispatchHook{
    24  		payload: alloc.Job.Payload,
    25  	}
    26  	h.logger = logger.Named(h.Name())
    27  	return h
    28  }
    29  
    30  func (*dispatchHook) Name() string {
    31  	// Copied in client/state when upgrading from <0.9 schemas, so if you
    32  	// change it here you also must change it there.
    33  	return "dispatch_payload"
    34  }
    35  
    36  func (h *dispatchHook) Prestart(ctx context.Context, req *interfaces.TaskPrestartRequest, resp *interfaces.TaskPrestartResponse) error {
    37  	if len(h.payload) == 0 || req.Task.DispatchPayload == nil || req.Task.DispatchPayload.File == "" {
    38  		// No dispatch payload
    39  		resp.Done = true
    40  		return nil
    41  	}
    42  
    43  	err := writeDispatchPayload(req.TaskDir.LocalDir, req.Task.DispatchPayload.File, h.payload)
    44  	if err != nil {
    45  		return err
    46  	}
    47  
    48  	h.logger.Trace("dispatch payload written",
    49  		"path", req.TaskDir.LocalDir,
    50  		"filename", req.Task.DispatchPayload.File,
    51  		"bytes", len(h.payload),
    52  	)
    53  
    54  	// Dispatch payload written successfully; mark as done
    55  	resp.Done = true
    56  	return nil
    57  }
    58  
    59  // writeDispatchPayload writes the payload to the given file or returns an
    60  // error.
    61  func writeDispatchPayload(base, filename string, payload []byte) error {
    62  	renderTo := filepath.Join(base, filename)
    63  	decoded, err := snappy.Decode(nil, payload)
    64  	if err != nil {
    65  		return err
    66  	}
    67  
    68  	if err := os.MkdirAll(filepath.Dir(renderTo), 0777); err != nil {
    69  		return err
    70  	}
    71  
    72  	return ioutil.WriteFile(renderTo, decoded, 0777)
    73  }