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

     1  package actor
     2  
     3  import (
     4  	"bufio"
     5  	"fmt"
     6  	"net"
     7  	"sync"
     8  
     9  	core "github.com/awesome-flow/flow/pkg/corev1alpha1"
    10  	"github.com/awesome-flow/flow/pkg/types"
    11  )
    12  
    13  type ReceiverUDP struct {
    14  	name   string
    15  	ctx    *core.Context
    16  	addr   *net.UDPAddr
    17  	conn   *net.UDPConn
    18  	queue  chan *core.Message
    19  	done   chan struct{}
    20  	wgconn sync.WaitGroup
    21  	wgpeer sync.WaitGroup
    22  }
    23  
    24  var _ core.Actor = (*ReceiverUDP)(nil)
    25  
    26  func NewReceiverUDP(name string, ctx *core.Context, params core.Params) (core.Actor, error) {
    27  	bind, ok := params["bind"]
    28  	if !ok {
    29  		return nil, fmt.Errorf("udp receiver is missing `bind` config")
    30  	}
    31  	addr, err := net.ResolveUDPAddr("udp", bind.(string))
    32  	if err != nil {
    33  		return nil, err
    34  	}
    35  
    36  	return &ReceiverUDP{
    37  		name:  name,
    38  		ctx:   ctx,
    39  		addr:  addr,
    40  		queue: make(chan *core.Message),
    41  		done:  make(chan struct{}),
    42  	}, nil
    43  }
    44  
    45  func (r *ReceiverUDP) Name() string {
    46  	return r.name
    47  }
    48  
    49  func (r *ReceiverUDP) handleConn(conn net.Conn) {
    50  	reader := bufio.NewReader(conn)
    51  	scanner := bufio.NewScanner(reader)
    52  
    53  	for scanner.Scan() {
    54  		msg := core.NewMessage(scanner.Bytes())
    55  		r.queue <- msg
    56  	}
    57  
    58  	if err := scanner.Err(); err != nil {
    59  		r.ctx.Logger().Error(err.Error())
    60  	}
    61  }
    62  
    63  func (r *ReceiverUDP) Start() error {
    64  	r.ctx.Logger().Info("starting udp listener at %s", r.addr)
    65  	conn, err := net.ListenUDP("udp", r.addr)
    66  	if err != nil {
    67  		return err
    68  	}
    69  	r.conn = conn
    70  
    71  	isdone := false
    72  	go func() {
    73  		<-r.done
    74  		isdone = true
    75  	}()
    76  
    77  	nthreads, ok := r.ctx.Config().Get(types.NewKey("system.maxprocs"))
    78  	if !ok {
    79  		nthreads = 1
    80  	}
    81  
    82  	for i := 0; i < nthreads.(int); i++ {
    83  		r.wgconn.Add(1)
    84  		go func() {
    85  			for !isdone {
    86  				r.handleConn(conn)
    87  			}
    88  
    89  			r.conn.Close()
    90  			r.wgconn.Done()
    91  		}()
    92  	}
    93  
    94  	return nil
    95  }
    96  
    97  func (r *ReceiverUDP) Stop() error {
    98  	close(r.done)
    99  	r.wgconn.Wait()
   100  	close(r.queue)
   101  	r.wgpeer.Wait()
   102  
   103  	return nil
   104  }
   105  
   106  func (r *ReceiverUDP) Connect(nthreads int, peer core.Receiver) error {
   107  	for i := 0; i < nthreads; i++ {
   108  		r.wgpeer.Add(1)
   109  		go func() {
   110  			for msg := range r.queue {
   111  				if err := peer.Receive(msg); err != nil {
   112  					r.ctx.Logger().Error(err.Error())
   113  				}
   114  			}
   115  			r.wgpeer.Done()
   116  		}()
   117  	}
   118  	return nil
   119  }
   120  
   121  func (u *ReceiverUDP) Receive(*core.Message) error {
   122  	return fmt.Errorf("udp receiver %q can not receive internal messages", u.name)
   123  }