github.com/ergo-services/ergo@v1.999.224/gen/pool.go (about)

     1  package gen
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"github.com/ergo-services/ergo/etf"
     7  	"github.com/ergo-services/ergo/lib"
     8  )
     9  
    10  type PoolBehavior interface {
    11  	ServerBehavior
    12  
    13  	InitPool(process *PoolProcess, args ...etf.Term) (PoolOptions, error)
    14  }
    15  
    16  type PoolProcess struct {
    17  	ServerProcess
    18  	options  PoolOptions
    19  	workers  []etf.Pid
    20  	monitors map[etf.Ref]int
    21  	i        int
    22  }
    23  
    24  type Pool struct {
    25  	Server
    26  }
    27  
    28  type PoolOptions struct {
    29  	NumWorkers    int
    30  	Worker        PoolWorkerBehavior
    31  	WorkerOptions ProcessOptions
    32  	WorkerArgs    []etf.Term
    33  }
    34  
    35  func (p *Pool) Init(process *ServerProcess, args ...etf.Term) error {
    36  	behavior, ok := process.Behavior().(PoolBehavior)
    37  	if !ok {
    38  		return fmt.Errorf("Pool: not a PoolBehavior")
    39  	}
    40  
    41  	pool := &PoolProcess{
    42  		ServerProcess: *process,
    43  		monitors:      make(map[etf.Ref]int),
    44  	}
    45  
    46  	// do not inherit parent State
    47  	pool.State = nil
    48  	poolOptions, err := behavior.InitPool(pool, args...)
    49  	if err != nil {
    50  		return err
    51  	}
    52  
    53  	poolOptions.WorkerOptions.Context = process.Context()
    54  	pool.options = poolOptions
    55  	process.State = pool
    56  
    57  	for i := 0; i < poolOptions.NumWorkers; i++ {
    58  		w, err := process.Spawn("", poolOptions.WorkerOptions, poolOptions.Worker,
    59  			poolOptions.WorkerArgs...)
    60  		if err != nil {
    61  			return err
    62  		}
    63  
    64  		pool.workers = append(pool.workers, w.Self())
    65  		ref := process.MonitorProcess(w.Self())
    66  		pool.monitors[ref] = i
    67  	}
    68  
    69  	return nil
    70  }
    71  
    72  func (p *Pool) HandleCall(process *ServerProcess, from ServerFrom, message etf.Term) (etf.Term, ServerStatus) {
    73  	pool := process.State.(*PoolProcess)
    74  	msg := workerCallMessage{
    75  		from:    from,
    76  		message: message,
    77  	}
    78  	if err := p.send(pool, msg); err != nil {
    79  		lib.Warning("Pool (HandleCall): all workers are busy. Message dropped")
    80  	}
    81  	return nil, ServerStatusIgnore
    82  }
    83  func (p *Pool) HandleCast(process *ServerProcess, message etf.Term) ServerStatus {
    84  	pool := process.State.(*PoolProcess)
    85  	msg := workerCastMessage{
    86  		message: message,
    87  	}
    88  	if err := p.send(pool, msg); err != nil {
    89  		lib.Warning("Pool (HandleCast): all workers are busy. Message dropped")
    90  	}
    91  	return ServerStatusOK
    92  }
    93  func (p *Pool) HandleInfo(process *ServerProcess, message etf.Term) ServerStatus {
    94  	pool := process.State.(*PoolProcess)
    95  	switch m := message.(type) {
    96  	case MessageDown:
    97  		// worker terminated. restart it
    98  
    99  		i, exist := pool.monitors[m.Ref]
   100  		if exist == false {
   101  			break
   102  		}
   103  		delete(pool.monitors, m.Ref)
   104  		w, err := process.Spawn("", pool.options.WorkerOptions, pool.options.Worker,
   105  			pool.options.WorkerArgs...)
   106  		if err != nil {
   107  			panicMessage := fmt.Sprintf("Pool: can't restart worker - %s", err)
   108  			panic(panicMessage)
   109  		}
   110  		pool.workers[i] = w.Self()
   111  		return ServerStatusOK
   112  	}
   113  
   114  	if err := p.send(pool, message); err != nil {
   115  		lib.Warning("Pool (HandleInfo): all workers are busy. Message dropped")
   116  	}
   117  
   118  	return ServerStatusOK
   119  }
   120  
   121  func (p *Pool) send(pool *PoolProcess, message etf.Term) error {
   122  	for retry := 0; retry < pool.options.NumWorkers; retry++ {
   123  		pool.i++
   124  		worker := pool.workers[pool.i%pool.options.NumWorkers]
   125  		if err := pool.Send(worker, message); err == nil {
   126  			return ServerStatusOK
   127  		}
   128  	}
   129  
   130  	return fmt.Errorf("error")
   131  }