github.com/ncodes/nomad@v0.5.7-0.20170403112158-97adf4a74fb3/client/task_runner_plus.go (about)

     1  package client
     2  
     3  import (
     4  	"fmt"
     5  	"log"
     6  	"strconv"
     7  	"strings"
     8  	"time"
     9  
    10  	"os/exec"
    11  
    12  	"github.com/ncodes/cocoon/tools"
    13  )
    14  
    15  // TaskRunnerPlus includes unofficial task runner functionalities.
    16  // Functionalities include the ability to stop and remove a docker container
    17  // identified by an id specified in the task environment.
    18  type TaskRunnerPlus struct {
    19  	ContainerEnvKey   string
    20  	MemoryAllocEnvKey string
    21  	taskEnv           map[string]string
    22  	l                 *log.Logger
    23  }
    24  
    25  // NewTaskRunnerPlus creates a new task runner
    26  func NewTaskRunnerPlus(logger *log.Logger, taskEnv map[string]string) *TaskRunnerPlus {
    27  	return &TaskRunnerPlus{
    28  		ContainerEnvKey:   "CONTAINER_ID",
    29  		MemoryAllocEnvKey: "ALLOC_MEMORY",
    30  		taskEnv:           taskEnv,
    31  		l:                 logger,
    32  	}
    33  }
    34  
    35  // Unofficial Feature: Deletes any docker image with a matching id as the `CONTAINER_ID` in the TaskEnv
    36  func (r *TaskRunnerPlus) stopContainer() error {
    37  
    38  	containerID := r.taskEnv[r.ContainerEnvKey]
    39  	if containerID == "" {
    40  		return fmt.Errorf("container id unspecified. No container to kill")
    41  	}
    42  
    43  	r.l.Printf("[DEBUG] driver.raw_exec: Attempting to stop associated container (if any)")
    44  
    45  	err := tools.DeleteContainer(containerID, false, false, false)
    46  	if err != nil {
    47  		if err == tools.ErrContainerNotFound {
    48  			r.l.Printf("[DEBUG] driver.raw_exec: No associated container found")
    49  			return nil
    50  		}
    51  		return fmt.Errorf("failed to delete container attached to task")
    52  	}
    53  
    54  	r.l.Printf("[DEBUG] driver.raw_exec: Successfully stopped task container")
    55  	return nil
    56  }
    57  
    58  // SendGRPCSignal sends an artifical signal to a GRPC server running
    59  // on a container (or anywhere). If request does not return within the timeout
    60  // period it returns with no error
    61  func (r *TaskRunnerPlus) SendGRPCSignal(timeout time.Duration) error {
    62  	return nil // not implemented
    63  }
    64  
    65  // KillOnLowMemory will kill a task if an expected amount
    66  // of memory is unavailable.
    67  func (r *TaskRunnerPlus) KillOnLowMemory(requiredMemMB int, kill func() error) error {
    68  	availableMemStr, err := exec.Command("bash", "-c", "free -m | grep Mem | awk '{print $4}'").Output()
    69  	if err != nil {
    70  		return fmt.Errorf("failed to check available memory. %s", err)
    71  	}
    72  	avMem, _ := strconv.Atoi(strings.TrimSpace(string(availableMemStr)))
    73  	if avMem < requiredMemMB {
    74  		return kill()
    75  	}
    76  	return nil
    77  }