github.com/awesome-flow/flow@v0.0.3-0.20190918184116-508d75d68a2c/pkg/corev1alpha1/actor/router.go (about)

     1  package actor
     2  
     3  import (
     4  	"sync"
     5  
     6  	core "github.com/awesome-flow/flow/pkg/corev1alpha1"
     7  )
     8  
     9  type Router struct {
    10  	name  string
    11  	ctx   *core.Context
    12  	rtmap map[string]chan *core.Message
    13  	lock  sync.Mutex
    14  	wg    sync.WaitGroup
    15  }
    16  
    17  var _ core.Actor = (*Router)(nil)
    18  
    19  func NewRouter(name string, ctx *core.Context, params core.Params) (core.Actor, error) {
    20  	return &Router{
    21  		name:  name,
    22  		ctx:   ctx,
    23  		rtmap: make(map[string]chan *core.Message),
    24  		lock:  sync.Mutex{},
    25  	}, nil
    26  }
    27  
    28  func (r *Router) Name() string {
    29  	return r.name
    30  }
    31  
    32  func (r *Router) Start() error {
    33  	return nil
    34  }
    35  
    36  func (r *Router) Stop() error {
    37  	for _, ch := range r.rtmap {
    38  		close(ch)
    39  	}
    40  	r.wg.Wait()
    41  	return nil
    42  }
    43  
    44  func (r *Router) Connect(nthreads int, peer core.Receiver) error {
    45  	r.lock.Lock()
    46  	defer r.lock.Unlock()
    47  	peername := peer.(core.Namer).Name()
    48  	if _, ok := r.rtmap[peername]; !ok {
    49  		r.rtmap[peername] = make(chan *core.Message)
    50  	}
    51  	queue := r.rtmap[peername]
    52  	for i := 0; i < nthreads; i++ {
    53  		r.wg.Add(1)
    54  		go func() {
    55  			for msg := range queue {
    56  				if err := peer.Receive(msg); err != nil {
    57  					msg.Complete(core.MsgStatusFailed)
    58  					r.ctx.Logger().Error(err.Error())
    59  				}
    60  			}
    61  			r.wg.Done()
    62  		}()
    63  	}
    64  	return nil
    65  }
    66  
    67  func (r *Router) Receive(msg *core.Message) error {
    68  	if rtkey, ok := msg.Meta("sendto"); ok {
    69  		if queue, ok := r.rtmap[rtkey.(string)]; ok {
    70  			queue <- msg
    71  			return nil
    72  		}
    73  	}
    74  	msg.Complete(core.MsgStatusUnroutable)
    75  	return nil
    76  }