github.com/oarkflow/log@v1.0.78/gelf/tcpreader.go (about)

     1  package gelf
     2  
     3  import (
     4  	"bufio"
     5  	"encoding/json"
     6  	"fmt"
     7  	"net"
     8  	"time"
     9  )
    10  
    11  type TCPReader struct {
    12  	listener *net.TCPListener
    13  	conn     net.Conn
    14  	messages chan []byte
    15  }
    16  
    17  type connChannels struct {
    18  	drop    chan string
    19  	confirm chan string
    20  }
    21  
    22  func newTCPReader(addr string) (*TCPReader, chan string, chan string, error) {
    23  	var err error
    24  	tcpAddr, err := net.ResolveTCPAddr("tcp", addr)
    25  	if err != nil {
    26  		return nil, nil, nil, fmt.Errorf("ResolveTCPAddr('%s'): %s", addr, err)
    27  	}
    28  
    29  	listener, err := net.ListenTCP("tcp", tcpAddr)
    30  	if err != nil {
    31  		return nil, nil, nil, fmt.Errorf("ListenTCP: %s", err)
    32  	}
    33  
    34  	r := &TCPReader{
    35  		listener: listener,
    36  		messages: make(chan []byte, 100), // Make a buffered channel with at most 100 messages
    37  	}
    38  
    39  	closeSignal := make(chan string, 1)
    40  	doneSignal := make(chan string, 1)
    41  
    42  	go r.listenUntilCloseSignal(closeSignal, doneSignal)
    43  
    44  	return r, closeSignal, doneSignal, nil
    45  }
    46  
    47  func (r *TCPReader) accepter(connections chan net.Conn) {
    48  	for {
    49  		conn, err := r.listener.Accept()
    50  		if err != nil {
    51  			break
    52  		}
    53  		connections <- conn
    54  	}
    55  }
    56  
    57  func (r *TCPReader) listenUntilCloseSignal(closeSignal chan string, doneSignal chan string) {
    58  	defer func() { doneSignal <- "done" }()
    59  	defer r.listener.Close()
    60  	var conns []connChannels
    61  	connectionsChannel := make(chan net.Conn, 1)
    62  	go r.accepter(connectionsChannel)
    63  	for {
    64  		select {
    65  		case conn := <-connectionsChannel:
    66  			dropSignal := make(chan string, 1)
    67  			dropConfirm := make(chan string, 1)
    68  			channels := connChannels{drop: dropSignal, confirm: dropConfirm}
    69  			go handleConnection(conn, r.messages, dropSignal, dropConfirm)
    70  			conns = append(conns, channels)
    71  		default:
    72  		}
    73  
    74  		select {
    75  		case sig := <-closeSignal:
    76  			if sig == "stop" || sig == "drop" {
    77  				if len(conns) >= 1 {
    78  					for _, s := range conns {
    79  						if s.drop != nil {
    80  							s.drop <- "drop"
    81  							<-s.confirm
    82  							conns = append(conns[:0], conns[1:]...)
    83  						}
    84  					}
    85  					if sig == "stop" {
    86  						return
    87  					}
    88  				} else if sig == "stop" {
    89  					closeSignal <- "stop"
    90  				}
    91  				if sig == "drop" {
    92  					doneSignal <- "done"
    93  				}
    94  			}
    95  		default:
    96  		}
    97  	}
    98  }
    99  
   100  func (r *TCPReader) addr() string {
   101  	return r.listener.Addr().String()
   102  }
   103  
   104  func handleConnection(conn net.Conn, messages chan<- []byte, dropSignal chan string, dropConfirm chan string) {
   105  	defer func() { dropConfirm <- "done" }()
   106  	defer conn.Close()
   107  	reader := bufio.NewReader(conn)
   108  
   109  	var b []byte
   110  	var err error
   111  	drop := false
   112  	canDrop := false
   113  
   114  	for {
   115  		conn.SetDeadline(time.Now().Add(2 * time.Second))
   116  		if b, err = reader.ReadBytes(0); err != nil {
   117  			if drop {
   118  				return
   119  			}
   120  		} else if len(b) > 0 {
   121  			messages <- b
   122  			canDrop = true
   123  			if drop {
   124  				return
   125  			}
   126  		} else if drop {
   127  			return
   128  		}
   129  		select {
   130  		case sig := <-dropSignal:
   131  			if sig == "drop" {
   132  				drop = true
   133  				time.Sleep(1 * time.Second)
   134  				if canDrop {
   135  					return
   136  				}
   137  			}
   138  		default:
   139  		}
   140  	}
   141  }
   142  
   143  func (r *TCPReader) readMessage() (*Message, error) {
   144  	b := <-r.messages
   145  
   146  	var msg Message
   147  	if err := json.Unmarshal(b[:len(b)-1], &msg); err != nil {
   148  		return nil, fmt.Errorf("json.Unmarshal: %s", err)
   149  	}
   150  
   151  	return &msg, nil
   152  }
   153  
   154  func (r *TCPReader) Close() {
   155  	r.listener.Close()
   156  }