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

     1  package gelf
     2  
     3  import (
     4  	"fmt"
     5  	"net"
     6  	"os"
     7  	"sync"
     8  	"time"
     9  )
    10  
    11  const (
    12  	DefaultMaxReconnect   = 3
    13  	DefaultReconnectDelay = 1
    14  )
    15  
    16  type TCPWriter struct {
    17  	Writer
    18  	mu             sync.Mutex
    19  	MaxReconnect   int
    20  	ReconnectDelay time.Duration
    21  }
    22  
    23  func NewTCPWriter(addr string) (*TCPWriter, error) {
    24  	var err error
    25  	w := new(TCPWriter)
    26  	w.MaxReconnect = DefaultMaxReconnect
    27  	w.ReconnectDelay = DefaultReconnectDelay
    28  	w.proto = "tcp"
    29  	w.addr = addr
    30  
    31  	if w.conn, err = net.Dial("tcp", addr); err != nil {
    32  		return nil, err
    33  	}
    34  	if w.hostname, err = os.Hostname(); err != nil {
    35  		return nil, err
    36  	}
    37  
    38  	return w, nil
    39  }
    40  
    41  // WriteMessage sends the specified message to the GELF server
    42  // specified in the call to New().  It assumes all the fields are
    43  // filled out appropriately.  In general, clients will want to use
    44  // Write, rather than WriteMessage.
    45  func (w *TCPWriter) WriteMessage(m *Message) (err error) {
    46  	buf := newBuffer()
    47  	defer bufPool.Put(buf)
    48  	messageBytes, err := m.toBytes(buf)
    49  	if err != nil {
    50  		return err
    51  	}
    52  
    53  	messageBytes = append(messageBytes, 0)
    54  
    55  	n, err := w.writeToSocketWithReconnectAttempts(messageBytes)
    56  	if err != nil {
    57  		return err
    58  	}
    59  	if n != len(messageBytes) {
    60  		return fmt.Errorf("bad write (%d/%d)", n, len(messageBytes))
    61  	}
    62  
    63  	return nil
    64  }
    65  
    66  func (w *TCPWriter) Write(p []byte) (n int, err error) {
    67  	file, line := getCallerIgnoringLogMulti(1)
    68  
    69  	m := constructMessage(p, w.hostname, w.Facility, file, line)
    70  
    71  	if err = w.WriteMessage(m); err != nil {
    72  		return 0, err
    73  	}
    74  
    75  	return len(p), nil
    76  }
    77  
    78  func (w *TCPWriter) writeToSocketWithReconnectAttempts(zBytes []byte) (n int, err error) {
    79  	var errConn error
    80  	var i int
    81  
    82  	w.mu.Lock()
    83  	for i = 0; i <= w.MaxReconnect; i++ {
    84  		errConn = nil
    85  
    86  		if w.conn != nil {
    87  			n, err = w.conn.Write(zBytes)
    88  		} else {
    89  			err = fmt.Errorf("Connection was nil, will attempt reconnect")
    90  		}
    91  		if err != nil {
    92  			time.Sleep(w.ReconnectDelay * time.Second)
    93  			w.conn, errConn = net.Dial("tcp", w.addr)
    94  		} else {
    95  			break
    96  		}
    97  	}
    98  	w.mu.Unlock()
    99  
   100  	if i > w.MaxReconnect {
   101  		return 0, fmt.Errorf("Maximum reconnection attempts was reached; giving up")
   102  	}
   103  	if errConn != nil {
   104  		return 0, fmt.Errorf("Write Failed: %s\nReconnection failed: %s", err, errConn)
   105  	}
   106  	return n, nil
   107  }