github.com/brownsys/tracing-framework-go@v0.0.0-20161210174012-0542a62412fe/xtrace/internal/pubsub/client.go (about)

     1  package pubsub
     2  
     3  import (
     4  	"encoding/binary"
     5  	"fmt"
     6  	"io"
     7  	"net"
     8  	"os"
     9  	"sync"
    10  	"sync/atomic"
    11  )
    12  
    13  // A Client represents a connection to a pubsub server.
    14  // The zero value is not a valid Client.
    15  type Client struct {
    16  	messages chan message
    17  	quit     chan struct{}
    18  	closed   uint32
    19  }
    20  
    21  // NewClient creates a new connection to server.
    22  func NewClient(server string) (c *Client, err error) {
    23  	conn, err := net.Dial("tcp", server)
    24  	if err != nil {
    25  		return nil, err
    26  	}
    27  
    28  	c = &Client{
    29  		messages: make(chan message, 1024),
    30  		quit:     make(chan struct{}, 1),
    31  	}
    32  	go c.daemon(server, conn)
    33  	return c, nil
    34  }
    35  
    36  func (c *Client) daemon(server string, conn net.Conn) {
    37  	for {
    38  		var m message
    39  		select {
    40  		case <-c.quit:
    41  			conn.Close()
    42  			return
    43  		case m = <-c.messages:
    44  		}
    45  
    46  		for {
    47  			err := writeMessage(conn, m)
    48  			if err == nil {
    49  				if m.wg != nil {
    50  					m.wg.Done()
    51  				}
    52  				break
    53  			}
    54  			fmt.Fprintf(os.Stderr, "pubsub client error: %v\n", err)
    55  
    56  			for {
    57  				conn, err = net.Dial("tcp", server)
    58  				if err == nil {
    59  					break
    60  				}
    61  				fmt.Fprintf(os.Stderr, "pubsub client error: %v\n", err)
    62  			}
    63  		}
    64  	}
    65  }
    66  
    67  // Close closes the connection with the server.
    68  func (c *Client) Close() {
    69  	select {
    70  	case c.quit <- struct{}{}:
    71  	default:
    72  	}
    73  	atomic.StoreUint32(&c.closed, 1)
    74  }
    75  
    76  // Publish publishes msg on the given topic. Publish may block,
    77  // but it is not guaranteed that when Publish returns, the message
    78  // has been received by the server.
    79  func (c *Client) Publish(topic, msg []byte) {
    80  	if atomic.LoadUint32(&c.closed) == 1 {
    81  		panic("publish on closed client")
    82  	}
    83  
    84  	c.messages <- message{topic: topic, message: msg}
    85  }
    86  
    87  // PublishBlock is like Publish, except that it blocks until
    88  // the message has been written to the server. Note that this
    89  // does not guarantee receipt by the server.
    90  func (c *Client) PublishBlock(topic, msg []byte) {
    91  	if atomic.LoadUint32(&c.closed) == 1 {
    92  		panic("publish on closed client")
    93  	}
    94  
    95  	var wg sync.WaitGroup
    96  	wg.Add(1)
    97  	c.messages <- message{topic: topic, message: msg, wg: &wg}
    98  	wg.Wait()
    99  }
   100  
   101  // PublishString is equivalent to Publish([]byte(topic), []byte(msg)).
   102  func (c *Client) PublishString(topic, msg string) {
   103  	c.Publish([]byte(topic), []byte(msg))
   104  }
   105  
   106  // PublishStringBlock is equivalent to
   107  // PublishBlock([]byte(topic), []byte(msg)).
   108  func (c *Client) PublishStringBlock(topic, msg string) {
   109  	c.PublishBlock([]byte(topic), []byte(msg))
   110  }
   111  
   112  type message struct {
   113  	topic   []byte
   114  	message []byte
   115  	wg      *sync.WaitGroup // used by PublishBlock to block until sent
   116  }
   117  
   118  func writeMessage(w io.Writer, m message) error {
   119  	l := 8 + len(m.topic) + len(m.message)
   120  	buf := make([]byte, l)
   121  	binary.BigEndian.PutUint32(buf, uint32(len(m.topic)))
   122  	copy(buf[4:], m.topic)
   123  	binary.BigEndian.PutUint32(buf[4+len(m.topic):], uint32(len(m.message)))
   124  	copy(buf[8+len(m.topic):], m.message)
   125  
   126  	n, err := w.Write(buf)
   127  	if err != nil && n < len(buf) {
   128  		return err
   129  	}
   130  	return nil
   131  }