path: root/pkg/worker_watcher/worker_watcher.go

package worker_watcher //nolint:golint,stylecheck

import (
	"context"
	"runtime"
	"sync"
	"time"

	"github.com/spiral/errors"
	"github.com/spiral/roadrunner/v2/interfaces/events"
	"github.com/spiral/roadrunner/v2/interfaces/worker"
	"github.com/spiral/roadrunner/v2/internal"
)

type Stack struct {
	workers            []worker.BaseProcess
	mutex              sync.RWMutex
	destroy            bool
	actualNumOfWorkers int64
}

func NewWorkersStack() *Stack {
	w := runtime.NumCPU()
	return &Stack{
		workers:            make([]worker.BaseProcess, 0, w),
		actualNumOfWorkers: 0,
	}
}

func (stack *Stack) Reset() {
	stack.mutex.Lock()
	defer stack.mutex.Unlock()
	stack.actualNumOfWorkers = 0
	stack.workers = nil
}

// Push worker back to the stack
// If stack in destroy state, Push will provide 100ms window to unlock the mutex
func (stack *Stack) Push(w worker.BaseProcess) {
	stack.mutex.Lock()
	defer stack.mutex.Unlock()
	stack.actualNumOfWorkers++
	stack.workers = append(stack.workers, w)
}

func (stack *Stack) IsEmpty() bool {
	stack.mutex.Lock()
	defer stack.mutex.Unlock()
	return len(stack.workers) == 0
}

func (stack *Stack) Pop() (worker.BaseProcess, bool) {
	stack.mutex.Lock()
	defer stack.mutex.Unlock()

	// do not release new stack
	if stack.destroy {
		return nil, true
	}

	if len(stack.workers) == 0 {
		return nil, false
	}

	// move worker
	w := stack.workers[len(stack.workers)-1]
	stack.workers = stack.workers[:len(stack.workers)-1]
	stack.actualNumOfWorkers--
	return w, false
}

func (stack *Stack) FindAndRemoveByPid(pid int64) bool {
	stack.mutex.Lock()
	defer stack.mutex.Unlock()
	for i := 0; i < len(stack.workers); i++ {
		// worker in the stack, reallocating
		if stack.workers[i].Pid() == pid {
			stack.workers = append(stack.workers[:i], stack.workers[i+1:]...)
			stack.actualNumOfWorkers--
			// worker found and removed
			return true
		}
	}
	// no worker with such ID
	return false
}

// Workers return copy of the workers in the stack
func (stack *Stack) Workers() []worker.BaseProcess {
	stack.mutex.Lock()
	defer stack.mutex.Unlock()
	workersCopy := make([]worker.BaseProcess, 0, 1)
	// copy
	for _, v := range stack.workers {
		workersCopy = append(workersCopy, v)
	}

	return workersCopy
}

func (stack *Stack) isDestroying() bool {
	stack.mutex.Lock()
	defer stack.mutex.Unlock()
	return stack.destroy
}

// we also have to give a chance to pool to Push worker (return it)
func (stack *Stack) Destroy(ctx context.Context) {
	stack.mutex.Lock()
	stack.destroy = true
	stack.mutex.Unlock()

	tt := time.NewTicker(time.Millisecond * 100)
	for {
		select {
		case <-tt.C:
			stack.mutex.Lock()
			// that might be one of the workers is working
			if len(stack.workers) != int(stack.actualNumOfWorkers) {
				stack.mutex.Unlock()
				continue
			}
			stack.mutex.Unlock()
			// unnecessary mutex, but
			// just to make sure. All stack at this moment are in the stack
			// Pop operation is blocked, push can't be done, since it's not possible to pop
			stack.mutex.Lock()
			for i := 0; i < len(stack.workers); i++ {
				// set state for the stack in the stack (unused at the moment)
				stack.workers[i].State().Set(internal.StateDestroyed)
				// kill the worker
				_ = stack.workers[i].Kill()
			}
			stack.mutex.Unlock()
			tt.Stop()
			// clear
			stack.Reset()
			return
		}
	}
}

// workerCreateFunc can be nil, but in that case, dead stack will not be replaced
func NewWorkerWatcher(allocator worker.Allocator, numWorkers int64, events events.Handler) worker.Watcher {
	ww := &workerWatcher{
		stack:             NewWorkersStack(),
		allocator:         allocator,
		initialNumWorkers: numWorkers,
		actualNumWorkers:  numWorkers,
		events:            events,
	}

	return ww
}

type workerWatcher struct {
	mutex             sync.RWMutex
	stack             *Stack
	allocator         worker.Allocator
	initialNumWorkers int64
	actualNumWorkers  int64
	events            events.Handler
}

func (ww *workerWatcher) AddToWatch(workers []worker.BaseProcess) error {
	for i := 0; i < len(workers); i++ {
		ww.stack.Push(workers[i])

		go func(swc worker.BaseProcess) {
			ww.wait(swc)
		}(workers[i])
	}
	return nil
}

func (ww *workerWatcher) GetFreeWorker(ctx context.Context) (worker.BaseProcess, error) {
	const op = errors.Op("GetFreeWorker")
	// thread safe operation
	w, stop := ww.stack.Pop()
	if stop {
		return nil, errors.E(op, errors.WatcherStopped)
	}

	// handle worker remove state
	// in this state worker is destroyed by supervisor
	if w != nil && w.State().Value() == internal.StateRemove {
		err := ww.RemoveWorker(w)
		if err != nil {
			return nil, err
		}
		// try to get next
		return ww.GetFreeWorker(ctx)
	}
	// no free stack
	if w == nil {
		for {
			select {
			default:
				w, stop = ww.stack.Pop()
				if stop {
					return nil, errors.E(op, errors.WatcherStopped)
				}
				if w == nil {
					continue
				}
				return w, nil
			case <-ctx.Done():
				return nil, errors.E(op, errors.NoFreeWorkers, errors.Str("no free workers in the stack, timeout exceed"))
			}
		}
	}

	return w, nil
}

func (ww *workerWatcher) AllocateNew() error {
	ww.stack.mutex.Lock()
	const op = errors.Op("allocate new worker")
	sw, err := ww.allocator()
	if err != nil {
		return errors.E(op, errors.WorkerAllocate, err)
	}

	ww.addToWatch(sw)
	ww.stack.mutex.Unlock()
	ww.PushWorker(sw)

	return nil
}

func (ww *workerWatcher) RemoveWorker(wb worker.BaseProcess) error {
	ww.mutex.Lock()
	defer ww.mutex.Unlock()

	const op = errors.Op("remove worker")
	pid := wb.Pid()

	if ww.stack.FindAndRemoveByPid(pid) {
		wb.State().Set(internal.StateRemove)
		err := wb.Kill()
		if err != nil {
			return errors.E(op, err)
		}
		return nil
	}

	return nil
}

// O(1) operation
func (ww *workerWatcher) PushWorker(w worker.BaseProcess) {
	ww.mutex.Lock()
	defer ww.mutex.Unlock()
	ww.stack.Push(w)
}

// Destroy all underlying stack (but let them to complete the task)
func (ww *workerWatcher) Destroy(ctx context.Context) {
	// destroy stack, we don't use ww mutex here, since we should be able to push worker
	ww.stack.Destroy(ctx)
}

// Warning, this is O(n) operation, and it will return copy of the actual workers
func (ww *workerWatcher) WorkersList() []worker.BaseProcess {
	return ww.stack.Workers()
}

func (ww *workerWatcher) wait(w worker.BaseProcess) {
	const op = errors.Op("process wait")
	err := w.Wait()
	if err != nil {
		ww.events.Push(events.WorkerEvent{
			Event:   events.EventWorkerError,
			Worker:  w,
			Payload: errors.E(op, err),
		})
	}

	if w.State().Value() == internal.StateDestroyed {
		// worker was manually destroyed, no need to replace
		ww.events.Push(events.PoolEvent{Event: events.EventWorkerDestruct, Payload: w})
		return
	}

	_ = ww.stack.FindAndRemoveByPid(w.Pid())
	err = ww.AllocateNew()
	if err != nil {
		ww.events.Push(events.PoolEvent{
			Event:   events.EventPoolError,
			Payload: errors.E(op, err),
		})
	}
}

func (ww *workerWatcher) addToWatch(wb worker.BaseProcess) {
	go func() {
		ww.wait(wb)
	}()
}