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

     1  package actor
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"fmt"
     7  	"net"
     8  	"sync"
     9  	"time"
    10  
    11  	"github.com/awesome-flow/flow/pkg/cfg"
    12  	core "github.com/awesome-flow/flow/pkg/corev1alpha1"
    13  	"github.com/awesome-flow/flow/pkg/types"
    14  )
    15  
    16  const (
    17  	ConnReadTimeout  = 50 * time.Millisecond
    18  	ConnWriteTimeout = 50 * time.Millisecond
    19  	MsgSendTimeout   = 100 * time.Millisecond
    20  
    21  	DefaultBufSize = 4 * 1024
    22  )
    23  
    24  var (
    25  	TcpRespFail = []byte("FAILED\r\n")
    26  	TcpRespInvd = []byte("INVALID\r\n")
    27  	TcpRespPsnt = []byte("PARTSENT\r\n")
    28  	TcpRespOk   = []byte("OK\r\n")
    29  	TcpRespTime = []byte("TIMEOUT\r\n")
    30  	TcpRespThrt = []byte("THROTTLED\r\n")
    31  	TcpRespUnrt = []byte("UNROUTABLE\r\n")
    32  )
    33  
    34  var MsgStatusToTcpResp = map[core.MsgStatus][]byte{
    35  	core.MsgStatusDone:        TcpRespOk,
    36  	core.MsgStatusPartialSend: TcpRespPsnt,
    37  	core.MsgStatusInvalid:     TcpRespInvd,
    38  	core.MsgStatusFailed:      TcpRespFail,
    39  	core.MsgStatusTimedOut:    TcpRespTime,
    40  	core.MsgStatusUnroutable:  TcpRespUnrt,
    41  	core.MsgStatusThrottled:   TcpRespThrt,
    42  }
    43  
    44  type ReceiverTCP struct {
    45  	name     string
    46  	ctx      *core.Context
    47  	silent   bool
    48  	bufsize  int
    49  	addr     *net.TCPAddr
    50  	listener net.Listener
    51  	queue    chan *core.Message
    52  	done     chan struct{}
    53  	wgconn   sync.WaitGroup
    54  	wgpeer   sync.WaitGroup
    55  }
    56  
    57  var _ core.Actor = (*ReceiverTCP)(nil)
    58  
    59  func NewReceiverTCP(name string, ctx *core.Context, params core.Params) (core.Actor, error) {
    60  	bind, ok := params["bind"]
    61  	if !ok {
    62  		return nil, fmt.Errorf("tcp receiver %q is missing `bind` config", name)
    63  	}
    64  
    65  	addr, err := net.ResolveTCPAddr("tcp", bind.(string))
    66  	if err != nil {
    67  		return nil, err
    68  	}
    69  
    70  	var silent bool
    71  	if s, ok := params["silent"]; ok {
    72  		if s.(string) == "true" {
    73  			silent = true
    74  		} else if s.(string) != "false" {
    75  			return nil, fmt.Errorf("tcp receiver %q got an unexpected (non-bool) value for silent: %q", name, s)
    76  		}
    77  	}
    78  
    79  	bufsize, ok := params["buf_size"]
    80  	if !ok {
    81  		bufsize = DefaultBufSize
    82  	}
    83  
    84  	return &ReceiverTCP{
    85  		ctx:     ctx,
    86  		name:    name,
    87  		addr:    addr,
    88  		silent:  silent,
    89  		bufsize: bufsize.(int),
    90  		queue:   make(chan *core.Message),
    91  		done:    make(chan struct{}),
    92  	}, nil
    93  }
    94  
    95  func (r *ReceiverTCP) Name() string {
    96  	return r.name
    97  }
    98  
    99  func (r *ReceiverTCP) Start() error {
   100  	l, err := net.Listen("tcp", r.addr.String())
   101  	if err != nil {
   102  		return err
   103  	}
   104  	r.listener = l
   105  	nthreads, ok := r.ctx.Config().Get(types.NewKey(cfg.SystemMaxprocs))
   106  	if !ok {
   107  		return fmt.Errorf("failed to fetch %q config", cfg.SystemMaxprocs)
   108  	}
   109  
   110  	go func() {
   111  		<-r.done
   112  		r.ctx.Logger().Info("closing tcp listener at %s", r.addr.String())
   113  		if err := l.Close(); err != nil {
   114  			r.ctx.Logger().Error("failed to close tcp listener gracefuly: %s", err)
   115  		}
   116  	}()
   117  
   118  	for i := 0; i < nthreads.(int); i++ {
   119  		go func() {
   120  			r.ctx.Logger().Info("starting tcp listener at %s", r.addr.String())
   121  			for {
   122  				conn, err := l.Accept()
   123  				if err != nil {
   124  					r.ctx.Logger().Error(err.Error())
   125  					continue
   126  				}
   127  				go r.handleConn(conn)
   128  				select {
   129  				case <-r.done:
   130  					break
   131  				default:
   132  				}
   133  			}
   134  		}()
   135  	}
   136  
   137  	return nil
   138  }
   139  
   140  func (r *ReceiverTCP) Stop() error {
   141  	close(r.done)
   142  	r.wgconn.Wait()
   143  	close(r.queue)
   144  	r.wgpeer.Wait()
   145  	return nil
   146  }
   147  
   148  func (r *ReceiverTCP) Connect(nthreads int, peer core.Receiver) error {
   149  	for i := 0; i < nthreads; i++ {
   150  		r.wgpeer.Add(1)
   151  		go func() {
   152  			var err error
   153  			for msg := range r.queue {
   154  				if err = peer.Receive(msg); err != nil {
   155  					msg.Complete(core.MsgStatusFailed)
   156  					r.ctx.Logger().Error(err.Error())
   157  				}
   158  			}
   159  			r.wgpeer.Done()
   160  		}()
   161  	}
   162  	return nil
   163  }
   164  
   165  func (t *ReceiverTCP) Receive(*core.Message) error {
   166  	return fmt.Errorf("tcp receiver %q can not receive internal messages", t.name)
   167  }
   168  
   169  // dropCR drops a terminal \r from the data.
   170  func dropCR(data []byte) []byte {
   171  	if len(data) > 0 && data[len(data)-1] == '\r' {
   172  		return data[0 : len(data)-1]
   173  	}
   174  	return data
   175  }
   176  
   177  func ScanBin(data []byte, atEOF bool) (advance int, token []byte, err error) {
   178  	if atEOF && len(data) == 0 {
   179  		return 0, nil, nil
   180  	}
   181  	if i := bytes.Index(data, []byte{'\r', '\n'}); i >= 0 {
   182  		// We have a full newline-terminated line.
   183  		return i + 2, dropCR(data[0:i]), nil
   184  	}
   185  	// If we're at EOF, we have a final, non-terminated line. Return it.
   186  	if atEOF {
   187  		return len(data), dropCR(data), nil
   188  	}
   189  	// Request more data.
   190  	return 0, nil, nil
   191  }
   192  
   193  func (r *ReceiverTCP) handleConn(conn net.Conn) {
   194  	r.ctx.Logger().Debug("new tcp connection from %s", conn.RemoteAddr())
   195  
   196  	r.wgconn.Add(1)
   197  
   198  	reader := bufio.NewReader(conn)
   199  	scanner := bufio.NewScanner(reader)
   200  	buf := make([]byte, 1024)
   201  	scanner.Buffer(buf, r.bufsize)
   202  	scanner.Split(ScanBin)
   203  
   204  	isdone := false
   205  	scanover := make(chan struct{})
   206  	go func() {
   207  		select {
   208  		case <-r.done:
   209  			isdone = true
   210  		case <-scanover:
   211  		}
   212  	}()
   213  
   214  	for !isdone && scanner.Scan() {
   215  		msg := core.NewMessage(scanner.Bytes())
   216  		r.queue <- msg
   217  
   218  		if r.silent {
   219  			continue
   220  		}
   221  
   222  		var status core.MsgStatus
   223  
   224  		select {
   225  		case s := <-msg.AwaitChan():
   226  			status = s
   227  		case <-time.After(MsgSendTimeout):
   228  			status = core.MsgStatusTimedOut
   229  		}
   230  
   231  		reply := MsgStatusToTcpResp[status]
   232  		conn.SetWriteDeadline(time.Now().Add(ConnWriteTimeout))
   233  		if _, err := conn.Write(reply); err != nil {
   234  			r.ctx.Logger().Error(err.Error())
   235  		}
   236  	}
   237  	close(scanover)
   238  	if err := scanner.Err(); err != nil {
   239  		r.ctx.Logger().Error(err.Error())
   240  	}
   241  
   242  	r.wgconn.Done()
   243  
   244  	r.ctx.Logger().Debug("closing tcp connnection from %s", conn.RemoteAddr())
   245  }