lambda/supervisor/local_supervisor.go

// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

package supervisor

import (
	"context"
	"errors"
	"fmt"
	"os/exec"
	"runtime"
	"sync"
	"syscall"
	"time"

	log "github.com/sirupsen/logrus"
	"go.amzn.com/lambda/supervisor/model"
)

// typecheck interface compliance
var _ model.SupervisorClient = (*LocalSupervisor)(nil)

type process struct {
	// pid of the running process
	pid int
	// channel that can be use to block
	// while waiting on process termination.
	termination chan struct{}
}

type LocalSupervisor struct {
	events               chan model.Event
	processMapLock       sync.Mutex
	processMap           map[string]process
	freezeThawCycleStart time.Time

	RootPath string
}

func NewLocalSupervisor() *LocalSupervisor {
	return &LocalSupervisor{
		events:     make(chan model.Event),
		processMap: make(map[string]process),
		RootPath:   "/",
	}
}

func (*LocalSupervisor) Start(ctx context.Context, req *model.StartRequest) error {
	return nil
}
func (*LocalSupervisor) Configure(ctx context.Context, req *model.ConfigureRequest) error {
	return nil
}
func (*LocalSupervisor) Exit(ctx context.Context) {}

func (s *LocalSupervisor) Exec(ctx context.Context, req *model.ExecRequest) error {
	if req.Domain != "runtime" {
		log.Debug("Exec is a no op if domain != runtime")
		return nil
	}
	command := exec.Command(req.Path, req.Args...)

	if req.Env != nil {
		envStrings := make([]string, 0, len(*req.Env))
		for key, value := range *req.Env {
			envStrings = append(envStrings, key+"="+value)
		}
		command.Env = envStrings
	}

	if req.Cwd != nil && *req.Cwd != "" {
		command.Dir = *req.Cwd
	}

	if req.ExtraFiles != nil {
		command.ExtraFiles = *req.ExtraFiles
	}

	command.Stdout = req.StdoutWriter
	command.Stderr = req.StderrWriter

	command.SysProcAttr = &syscall.SysProcAttr{Setpgid: true}

	err := command.Start()

	if err != nil {
		return err
		// TODO Use supevisor specific error
	}

	pid := command.Process.Pid
	termination := make(chan struct{})
	s.processMapLock.Lock()
	s.processMap[req.Name] = process{
		pid:         pid,
		termination: termination,
	}
	s.processMapLock.Unlock()

	// The first freeze thaw cycle starts on Exec() at init time
	s.freezeThawCycleStart = time.Now()

	go func() {
		err = command.Wait()
		// close the termination channel to unblock whoever's blocked on
		// it (used to implement kill's blocking behaviour)
		close(termination)

		var cell int32
		var exitStatus *int32
		var signo *int32
		var exitErr *exec.ExitError

		if err == nil {
			exitStatus = &cell
		} else if errors.As(err, &exitErr) {
			if status, ok := exitErr.Sys().(syscall.WaitStatus); ok {
				if code := status.ExitStatus(); code >= 0 {
					cell = int32(code)
					exitStatus = &cell
				} else {
					cell = int32(status.Signal())
					signo = &cell
				}
			}
		}

		if signo == nil && exitStatus == nil {
			log.Error("Cannot convert process exit status to unix WaitStatus. This is unexpected. Assuming ExitStatus 1")
			cell = 1
			exitStatus = &cell
		}
		s.events <- model.Event{
			Time: uint64(time.Now().UnixMilli()),
			Event: model.EventData{
				Domain:     &req.Domain,
				Name:       &req.Name,
				Signo:      signo,
				ExitStatus: exitStatus,
			},
		}
	}()

	return nil
}

func kill(p process, name string, deadline time.Time) error {
	// kill should report success if the process terminated by the time
	//supervisor receives the request.
	select {
	// if this case is selected, the channel is closed,
	// which means the process is terminated
	case <-p.termination:
		log.Debugf("Process %s already terminated.", name)
		return nil
	default:
		log.Infof("Sending SIGKILL to %s(%d).", name, p.pid)
	}

	if (time.Since(deadline)) > 0 {
		return fmt.Errorf("invalid timeout while killing %s", name)
	}

	pgid, err := syscall.Getpgid(p.pid)

	if err == nil {
		// Negative pid sends signal to all in process group
		syscall.Kill(-pgid, syscall.SIGKILL)
	} else {
		syscall.Kill(p.pid, syscall.SIGKILL)
	}

	ctx, cancel := context.WithDeadline(context.Background(), deadline)
	defer cancel()

	// block until the (main) process exits
	// or the timeout fires
	select {
	case <-p.termination:
		return nil
	case <-ctx.Done():
		return fmt.Errorf("timed out while trying to SIGKILL %s", name)
	}
}

func (s *LocalSupervisor) Kill(ctx context.Context, req *model.KillRequest) error {
	if req.Domain != "runtime" {
		log.Debug("Kill is a no op if domain != runtime")
		return nil
	}
	s.processMapLock.Lock()
	process, ok := s.processMap[req.Name]
	s.processMapLock.Unlock()
	if !ok {
		msg := "Unknown process"
		return &model.SupervisorError{
			Kind:    model.NoSuchEntity,
			Message: &msg,
		}
	}

	return kill(process, req.Name, req.Deadline)
}

func (s *LocalSupervisor) Terminate(ctx context.Context, req *model.TerminateRequest) error {
	if req.Domain != "runtime" {
		log.Debug("Terminate is no op if domain != runtime")
		return nil
	}
	s.processMapLock.Lock()
	process, ok := s.processMap[req.Name]
	pid := process.pid
	s.processMapLock.Unlock()
	if !ok {
		msg := "Unknown process"
		err := &model.SupervisorError{
			Kind:    model.NoSuchEntity,
			Message: &msg,
		}
		log.WithError(err).Errorf("Process %s not found in local supervisor map", req.Name)
		return err
	}

	pgid, err := syscall.Getpgid(pid)

	if err == nil {
		// Negative pid sends signal to all in process group
		// best effort, ignore errors
		_ = syscall.Kill(-pgid, syscall.SIGTERM)
	} else {
		_ = syscall.Kill(pid, syscall.SIGTERM)
	}

	return nil
}

func (s *LocalSupervisor) Stop(ctx context.Context, req *model.StopRequest) (*model.StopResponse, error) {
	if req.Domain != "runtime" {
		log.Debug("Shutdown is no op if domain != runtime")
		return &model.StopResponse{}, nil
	}

	// shut down kills all the processes in the map
	s.processMapLock.Lock()
	defer s.processMapLock.Unlock()

	nprocs := len(s.processMap)

	successes := make(chan struct{})
	errors := make(chan error)
	for name, proc := range s.processMap {
		go func(n string, p process) {
			log.Debugf("Killing %s", n)
			err := kill(p, n, req.Deadline)
			if err != nil {
				errors <- err
			} else {
				successes <- struct{}{}
			}

		}(name, proc)
	}

	var err error
	for i := 0; i < nprocs; i++ {
		select {
		case <-successes:
		case e := <-errors:
			if err == nil {
				err = fmt.Errorf("shutdown failed: %s", e.Error())
			}
		}

	}

	s.processMap = make(map[string]process)
	return nil, err
}

func (s *LocalSupervisor) Freeze(ctx context.Context, req *model.FreezeRequest) (*model.FreezeResponse, error) {
	// We return mocked freeze/thaw cycle metrics to mimic usage metrics in standalone mode
	var m runtime.MemStats
	runtime.ReadMemStats(&m)
	return &model.FreezeResponse{
		CycleDeltaMetrics: model.CycleDeltaMetrics{
			DomainCPURunNs:            uint64(time.Since(s.freezeThawCycleStart).Nanoseconds()),
			DomainRunNs:               uint64(time.Since(s.freezeThawCycleStart).Nanoseconds()),
			DomainMaxMemoryUsageBytes: m.Alloc,
			MicrovmCPURunNs:           uint64(time.Since(s.freezeThawCycleStart).Nanoseconds()),
		},
	}, nil
}
func (s *LocalSupervisor) Thaw(ctx context.Context, req *model.ThawRequest) error {
	s.freezeThawCycleStart = time.Now()
	return nil
}
func (s *LocalSupervisor) Ping(ctx context.Context) error {
	return nil
}

func (s *LocalSupervisor) Events(ctx context.Context, req *model.EventsRequest) (<-chan model.Event, error) {
	return s.events, nil
}