moira

1
package connection
2

3
import (
4
	"bufio"
5
	"io"
6
	"net"
7
	"sync"
8

9
	"go.avito.ru/DO/moira"
10
)
11

12
// Handler handling connection data and shift it to lineChan channel
13
type Handler struct {
14
	logger    moira.Logger
15
	wg        sync.WaitGroup
16
	terminate chan bool
17
}
18

19
// NewConnectionsHandler creates new Handler
20
func NewConnectionsHandler(logger moira.Logger) *Handler {
21
	return &Handler{
22
		logger:    logger,
23
		terminate: make(chan bool, 1),
24
	}
25
}
26

27
// HandleConnection convert every line from connection to metric and send it to lineChan channel
28
func (handler *Handler) HandleConnection(connection net.Conn, lineChan chan<- []byte) {
29
	handler.wg.Add(1)
30
	go func() {
31
		defer handler.wg.Done()
32
		handler.handle(connection, lineChan)
33
	}()
34
}
35

36
func (handler *Handler) handle(connection net.Conn, lineChan chan<- []byte) {
37
	buffer := bufio.NewReader(connection)
38

39
	go func(conn net.Conn) {
40
		<-handler.terminate
41
		conn.Close()
42
	}(connection)
43

44
	for {
45
		lineBytes, err := buffer.ReadBytes('\n')
46
		if err != nil {
47
			connection.Close()
48
			if err != io.EOF {
49
				handler.logger.ErrorF("read failed: %s", err)
50
			}
51
			break
52
		}
53
		lineBytes = lineBytes[:len(lineBytes)-1]
54
		lineChan <- lineBytes
55
	}
56
}
57

58
// StopHandlingConnections closes all open connections and wait for handling ramaining metrics
59
func (handler *Handler) StopHandlingConnections() {
60
	close(handler.terminate)
61
	handler.wg.Wait()
62
}
63