<commit_before>package client

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// allocSyncRetryIntv is the interval on which we retry updating
	// the status of the allocation
	allocSyncRetryIntv = 15 * time.Second
)

// taskStatus is used to track the status of a task
type taskStatus struct {
	Status      string
	Description string
}

// AllocStateUpdater is used to update the status of an allocation
type AllocStateUpdater func(alloc *structs.Allocation) error

// AllocRunner is used to wrap an allocation and provide the execution context.
type AllocRunner struct {
	config  *config.Config
	updater AllocStateUpdater
	logger  *log.Logger

	alloc *structs.Allocation

	dirtyCh chan struct{}

	ctx           *driver.ExecContext
	tasks         map[string]*TaskRunner
	RestartPolicy *structs.RestartPolicy
	taskLock      sync.RWMutex

	taskStatus     map[string]taskStatus
	taskStatusLock sync.RWMutex

	updateCh chan *structs.Allocation

	destroy     bool
	destroyCh   chan struct{}
	destroyLock sync.Mutex
	waitCh      chan struct{}
}

// allocRunnerState is used to snapshot the state of the alloc runner
type allocRunnerState struct {
	Alloc         *structs.Allocation
	RestartPolicy *structs.RestartPolicy
	TaskStatus    map[string]taskStatus
	Context       *driver.ExecContext
}

// NewAllocRunner is used to create a new allocation context
func NewAllocRunner(logger *log.Logger, config *config.Config, updater AllocStateUpdater, alloc *structs.Allocation) *AllocRunner {
	ar := &AllocRunner{
		config:     config,
		updater:    updater,
		logger:     logger,
		alloc:      alloc,
		dirtyCh:    make(chan struct{}, 1),
		tasks:      make(map[string]*TaskRunner),
		taskStatus: make(map[string]taskStatus),
		updateCh:   make(chan *structs.Allocation, 8),
		destroyCh:  make(chan struct{}),
		waitCh:     make(chan struct{}),
	}
	return ar
}

// stateFilePath returns the path to our state file
func (r *AllocRunner) stateFilePath() string {
	return filepath.Join(r.config.StateDir, "alloc", r.alloc.ID, "state.json")
}

// RestoreState is used to restore the state of the alloc runner
func (r *AllocRunner) RestoreState() error {
	// Load the snapshot
	var snap allocRunnerState
	if err := restoreState(r.stateFilePath(), &snap); err != nil {
		return err
	}

	// Restore fields
	r.alloc = snap.Alloc
	r.RestartPolicy = snap.RestartPolicy
	r.taskStatus = snap.TaskStatus
	r.ctx = snap.Context

	// Restore the task runners
	var mErr multierror.Error
	for name := range r.taskStatus {
		task := &structs.Task{Name: name}
		restartTracker := newRestartTracker(r.alloc.Job.Type, r.RestartPolicy)
		tr := NewTaskRunner(r.logger, r.config, r.setTaskStatus, r.ctx, r.alloc.ID, task, restartTracker)
		r.tasks[name] = tr
		if err := tr.RestoreState(); err != nil {
			r.logger.Printf("[ERR] client: failed to restore state for alloc %s task '%s': %v", r.alloc.ID, name, err)
			mErr.Errors = append(mErr.Errors, err)
		} else {
			go tr.Run()
		}
	}
	return mErr.ErrorOrNil()
}

// SaveState is used to snapshot our state
func (r *AllocRunner) SaveState(taskName string) error {
	r.taskStatusLock.RLock()
	snap := allocRunnerState{
		Alloc:         r.alloc,
		RestartPolicy: r.RestartPolicy,
		TaskStatus:    r.taskStatus,
		Context:       r.ctx,
	}
	err := persistState(r.stateFilePath(), &snap)
	r.taskStatusLock.RUnlock()
	if err != nil {
		return err
	}

	// Save state for each task
	r.taskLock.RLock()
	defer r.taskLock.RUnlock()
	var mErr multierror.Error
	if taskName != "" {
		tr, ok := r.tasks[taskName]
		if !ok {
			mErr.Errors = append(mErr.Errors, fmt.Errorf("[ERR] client: Task with name %v not found in alloc runner %v", taskName, r.alloc.Name))
		}
		r.saveTaskRunnerState(tr, &mErr)
		return mErr.ErrorOrNil()
	}
	for _, tr := range r.tasks {
		r.saveTaskRunnerState(tr, &mErr)
	}
	return mErr.ErrorOrNil()
}

func (r *AllocRunner) saveTaskRunnerState(tr *TaskRunner, mErr *multierror.Error) {
	if err := tr.SaveState(); err != nil {
		r.logger.Printf("[ERR] client: failed to save state for alloc %s task '%s': %v",
			r.alloc.ID, tr.task.Name, err)
		mErr.Errors = append(mErr.Errors, err)
	}
}

// DestroyState is used to cleanup after ourselves
func (r *AllocRunner) DestroyState() error {
	return os.RemoveAll(filepath.Dir(r.stateFilePath()))
}

// DestroyContext is used to destroy the context
func (r *AllocRunner) DestroyContext() error {
	return r.ctx.AllocDir.Destroy()
}
<commit_msg>Returning once we have an error while saving a task that doesn't exist<commit_after>package client

import (
	"encoding/json"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/hashicorp/go-multierror"
	"github.com/hashicorp/nomad/client/allocdir"
	"github.com/hashicorp/nomad/client/config"
	"github.com/hashicorp/nomad/client/driver"
	"github.com/hashicorp/nomad/nomad/structs"
)

const (
	// allocSyncRetryIntv is the interval on which we retry updating
	// the status of the allocation
	allocSyncRetryIntv = 15 * time.Second
)

// taskStatus is used to track the status of a task
type taskStatus struct {
	Status      string
	Description string
}

// AllocStateUpdater is used to update the status of an allocation
type AllocStateUpdater func(alloc *structs.Allocation) error

// AllocRunner is used to wrap an allocation and provide the execution context.
type AllocRunner struct {
	config  *config.Config
	updater AllocStateUpdater
	logger  *log.Logger

	alloc *structs.Allocation

	dirtyCh chan struct{}

	ctx           *driver.ExecContext
	tasks         map[string]*TaskRunner
	RestartPolicy *structs.RestartPolicy
	taskLock      sync.RWMutex

	taskStatus     map[string]taskStatus
	taskStatusLock sync.RWMutex

	updateCh chan *structs.Allocation

	destroy     bool
	destroyCh   chan struct{}
	destroyLock sync.Mutex
	waitCh      chan struct{}
}

// allocRunnerState is used to snapshot the state of the alloc runner
type allocRunnerState struct {
	Alloc         *structs.Allocation
	RestartPolicy *structs.RestartPolicy
	TaskStatus    map[string]taskStatus
	Context       *driver.ExecContext
}

// NewAllocRunner is used to create a new allocation context
func NewAllocRunner(logger *log.Logger, config *config.Config, updater AllocStateUpdater, alloc *structs.Allocation) *AllocRunner {
	ar := &AllocRunner{
		config:     config,
		updater:    updater,
		logger:     logger,
		alloc:      alloc,
		dirtyCh:    make(chan struct{}, 1),
		tasks:      make(map[string]*TaskRunner),
		taskStatus: make(map[string]taskStatus),
		updateCh:   make(chan *structs.Allocation, 8),
		destroyCh:  make(chan struct{}),
		waitCh:     make(chan struct{}),
	}
	return ar
}

// stateFilePath returns the path to our state file
func (r *AllocRunner) stateFilePath() string {
	return filepath.Join(r.config.StateDir, "alloc", r.alloc.ID, "state.json")
}

// RestoreState is used to restore the state of the alloc runner
func (r *AllocRunner) RestoreState() error {
	// Load the snapshot
	var snap allocRunnerState
	if err := restoreState(r.stateFilePath(), &snap); err != nil {
		return err
	}

	// Restore fields
	r.alloc = snap.Alloc
	r.RestartPolicy = snap.RestartPolicy
	r.taskStatus = snap.TaskStatus
	r.ctx = snap.Context

	// Restore the task runners
	var mErr multierror.Error
	for name := range r.taskStatus {
		task := &structs.Task{Name: name}
		restartTracker := newRestartTracker(r.alloc.Job.Type, r.RestartPolicy)
		tr := NewTaskRunner(r.logger, r.config, r.setTaskStatus, r.ctx, r.alloc.ID, task, restartTracker)
		r.tasks[name] = tr
		if err := tr.RestoreState(); err != nil {
			r.logger.Printf("[ERR] client: failed to restore state for alloc %s task '%s': %v", r.alloc.ID, name, err)
			mErr.Errors = append(mErr.Errors, err)
		} else {
			go tr.Run()
		}
	}
	return mErr.ErrorOrNil()
}

// SaveState is used to snapshot our state
func (r *AllocRunner) SaveState(taskName string) error {
	r.taskStatusLock.RLock()
	snap := allocRunnerState{
		Alloc:         r.alloc,
		RestartPolicy: r.RestartPolicy,
		TaskStatus:    r.taskStatus,
		Context:       r.ctx,
	}
	err := persistState(r.stateFilePath(), &snap)
	r.taskStatusLock.RUnlock()
	if err != nil {
		return err
	}

	// Save state for each task
	r.taskLock.RLock()
	defer r.taskLock.RUnlock()
	var mErr multierror.Error
	if taskName != "" {
		if tr, ok := r.tasks[taskName]; ok {
			r.saveTaskRunnerState(tr, &mErr)
		} else {
			mErr.Errors = append(mErr.Errors, fmt.Errorf("[ERR] client: Task with name %v not found in alloc runner %v", taskName, r.alloc.Name))
		}
		return mErr.ErrorOrNil()
	}
	for _, tr := range r.tasks {
		r.saveTaskRunnerState(tr, &mErr)
	}
	return mErr.ErrorOrNil()
}

func (r *AllocRunner) saveTaskRunnerState(tr *TaskRunner, mErr *multierror.Error) {
	if err := tr.SaveState(); err != nil {
		r.logger.Printf("[ERR] client: failed to save state for alloc %s task '%s': %v",
			r.alloc.ID, tr.task.Name, err)
		mErr.Errors = append(mErr.Errors, err)
	}
}

// DestroyState is used to cleanup after ourselves
func (r *AllocRunner) DestroyState() error {
	return os.RemoveAll(filepath.Dir(r.stateFilePath()))
}

// DestroyContext is used to destroy the context
func (r *AllocRunner) DestroyContext() error {
	return r.ctx.AllocDir.Destroy()
}

{'fraction_non_alphanumeric': 0.11304532427555461, 'fraction_numerical': 0.000849166755121537, 'mean_word_length': 3.82931829830856, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3911363', 'n_tokens_mistral': 3267, 'n_tokens_neox': 3180, 'n_words': 887}