github.com/artisanhe/tools@v1.0.1-0.20210607022958-19a8fef2eb04/catgo/cat-go/cat/sender.go (about)

     1  package cat
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/binary"
     6  	"net"
     7  	"time"
     8  
     9  	"github.com/artisanhe/tools/catgo/cat-go/message"
    10  )
    11  
    12  func createHeader() *message.Header {
    13  	return &message.Header{
    14  		Domain:   config.domain,
    15  		Hostname: config.hostname,
    16  		Ip:       config.ip,
    17  
    18  		MessageId:       manager.nextId(),
    19  		ParentMessageId: "",
    20  		RootMessageId:   "",
    21  	}
    22  }
    23  
    24  type catMessageSender struct {
    25  	scheduleMixin
    26  
    27  	normal  chan message.Messager
    28  	high    chan message.Messager
    29  	chConn  chan net.Conn
    30  	encoder message.Encoder
    31  
    32  	buf *bytes.Buffer
    33  
    34  	conn net.Conn
    35  }
    36  
    37  func (s *catMessageSender) GetName() string {
    38  	return "Sender"
    39  }
    40  
    41  func (s *catMessageSender) send(m message.Messager) {
    42  	var buf = s.buf
    43  	buf.Reset()
    44  
    45  	var header = createHeader()
    46  	if err := s.encoder.EncodeHeader(buf, header); err != nil {
    47  		return
    48  	}
    49  	if err := s.encoder.EncodeMessage(buf, m); err != nil {
    50  		return
    51  	}
    52  
    53  	var b = make([]byte, 4)
    54  	binary.BigEndian.PutUint32(b, uint32(buf.Len()))
    55  
    56  	if err := s.conn.SetWriteDeadline(time.Now().Add(time.Second * 3)); err != nil {
    57  		logger.Warning("Error occurred while setting write deadline, connection has been dropped.")
    58  		s.conn = nil
    59  		router.signals <- signalResetConnection
    60  	}
    61  	if _, err := s.conn.Write(b); err != nil {
    62  		logger.Warning("Error occurred while writing data, connection has been dropped.")
    63  		s.conn = nil
    64  		router.signals <- signalResetConnection
    65  		return
    66  	}
    67  	// fmt.Println(string(buf.Bytes()))
    68  	if _, err := s.conn.Write(buf.Bytes()); err != nil {
    69  		logger.Warning("Error occurred while writing data, connection has been dropped.")
    70  		s.conn = nil
    71  		router.signals <- signalResetConnection
    72  		return
    73  	}
    74  	return
    75  }
    76  
    77  func (s *catMessageSender) handleTransaction(trans *message.Transaction) {
    78  	if trans.GetStatus() != SUCCESS {
    79  		select {
    80  		case s.high <- trans:
    81  		default:
    82  			logger.Warning("High priority channel is full, transaction has been discarded.")
    83  		}
    84  	} else {
    85  		select {
    86  		case s.normal <- trans:
    87  		default:
    88  			// logger.Warning("Normal priority channel is full, transaction has been discarded.")
    89  		}
    90  	}
    91  }
    92  
    93  func (s *catMessageSender) handleEvent(event *message.Event) {
    94  	select {
    95  	case s.normal <- event:
    96  	default:
    97  		// logger.Warning("Normal priority channel is full, event has been discarded.")
    98  	}
    99  }
   100  
   101  func (s *catMessageSender) beforeStop() {
   102  	close(s.chConn)
   103  	close(s.high)
   104  	close(s.normal)
   105  
   106  	for m := range s.high {
   107  		s.send(m)
   108  	}
   109  	for m := range s.normal {
   110  		s.send(m)
   111  	}
   112  }
   113  
   114  func (s *catMessageSender) process() {
   115  	if s.conn == nil {
   116  		s.conn = <-s.chConn
   117  		logger.Info("Received a new connection: %s", s.conn.RemoteAddr().String())
   118  		return
   119  	}
   120  
   121  	select {
   122  	case sig := <-s.signals:
   123  		s.handle(sig)
   124  	case conn := <-s.chConn:
   125  		logger.Info("Received a new connection: %s", conn.RemoteAddr().String())
   126  		s.conn = conn
   127  	case m := <-s.high:
   128  		// logger.Debug("Receive a message [%s|%s] from high priority channel", m.GetType(), m.GetName())
   129  		s.send(m)
   130  	case m := <-s.normal:
   131  		// logger.Debug("Receive a message [%s|%s] from normal priority channel", m.GetType(), m.GetName())
   132  		s.send(m)
   133  	}
   134  }
   135  
   136  var sender = catMessageSender{
   137  	scheduleMixin: makeScheduleMixedIn(signalSenderExit),
   138  	normal:        make(chan message.Messager, normalPriorityQueueSize),
   139  	high:          make(chan message.Messager, highPriorityQueueSize),
   140  	chConn:        make(chan net.Conn),
   141  	encoder:       message.NewBinaryEncoder(),
   142  	buf:           bytes.NewBuffer([]byte{}),
   143  	conn:          nil,
   144  }