github.com/whiteboxio/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 }