sdk/workers/supervisor.go

/*
------------------------------------------------------------------------------------------------------------------------
####### sdk ####### Copyright (c) 2021-2022 losyme ################################################# MIT License #######
------------------------------------------------------------------------------------------------------------------------
*/

package workers

import (
	"forge.chapril.org/losyme/zombie"
)

type supervisor struct {
	poolSize int
	logger   Logger
	pool     *pool
	eventCh  chan string
	zombie   *zombie.Group
}

func newSupervisor(poolSize int, logger Logger, pool *pool) *supervisor {
	return &supervisor{
		poolSize: poolSize,
		logger:   logger,
		pool:     pool,
	}
}

func (s *supervisor) loop() {
	end := false

	for {
		if !end {
			if s.pool.size() < s.poolSize {
				s.pool.startOneWorker(s.eventCh)
				continue
			}

			if s.pool.size() > s.poolSize {
				s.pool.stopOneWorker()
				continue
			}
		}

		e := <-s.eventCh

		switch e {
		case "stop":
			if !end {
				end = true
				s.pool.stop()
			}
		case "0":
			if !end {
				s.poolSize = 0
			}
		case "-1":
			if !end && s.poolSize > 0 {
				s.poolSize--
			}
		case "+1":
			if !end && s.poolSize < _maxPoolSize {
				s.poolSize++
			}
		default: // wID
			size := s.pool.workerStopped(e, end)
			if end && size == 0 {
				return
			}
		}
	}
}

func (s *supervisor) start() {
	s.eventCh = make(chan string, 1)

	s.zombie = zombie.Go(
		1,
		func(_ int) {
			s.loop()
		},
		zombie.WithName("supervisor"),
		zombie.WithLogger(s.logger),
	)
}

func (s *supervisor) addOneWorker() {
	s.eventCh <- "+1"
}

func (s *supervisor) stopOneWorker() {
	s.eventCh <- "-1"
}

func (s *supervisor) stopAllWorkers() {
	s.eventCh <- "0"
}

func (s *supervisor) stop() {
	s.eventCh <- "stop"

	s.pool.wait()
	s.zombie.Wait()

	close(s.eventCh)
}

/*
######################################################################################################## @(°_°)@ #######
*/