influxdb/test/messaging.go

package test

import (
	"net/url"
	"sync"
	"time"

	"github.com/influxdb/influxdb"
	"github.com/influxdb/influxdb/messaging"
)

func init() {
	// Ensure the broker matches the handler's interface.
	_ = messaging.Handler{Broker: messaging.NewBroker()}
}

// MessagingClient represents a test client for the messaging broker.
type MessagingClient struct {
	mu      sync.Mutex
	index   uint64           // highest index
	conns   []*MessagingConn // list of all connections
	dataURL url.URL          // clients data node URL

	messagesByTopicID map[uint64][]*messaging.Message // message by topic

	PublishFunc func(*messaging.Message) (uint64, error)
	ConnFunc    func(topicID uint64) influxdb.MessagingConn
}

// NewMessagingClient returns a new instance of MessagingClient.
func NewMessagingClient(dataURL url.URL) *MessagingClient {
	c := &MessagingClient{
		messagesByTopicID: make(map[uint64][]*messaging.Message),
		dataURL:           dataURL,
	}
	c.PublishFunc = c.DefaultPublishFunc
	c.ConnFunc = c.DefaultConnFunc
	return c
}

// NewMessagingClient returns a new instance of MessagingClient.
func NewDefaultMessagingClient() *MessagingClient {
	testDataURL, _ := url.Parse("http://localhost:1234/data")
	return NewMessagingClient(*testDataURL)
}

func (c *MessagingClient) Open(path string) error { return nil }

// Close closes all open connections.
func (c *MessagingClient) Close() error {
	c.mu.Lock()
	defer c.mu.Unlock()

	for _, conn := range c.conns {
		conn.Close()
	}
	c.conns = nil

	return nil
}

func (c *MessagingClient) URLs() []url.URL   { return []url.URL{{Host: "local"}} }
func (c *MessagingClient) SetURLs([]url.URL) {}

func (c *MessagingClient) Publish(m *messaging.Message) (uint64, error) { return c.PublishFunc(m) }

// DefaultPublishFunc sets an autoincrementing index on the message and sends it to each topic connection.
func (c *MessagingClient) DefaultPublishFunc(m *messaging.Message) (uint64, error) {
	c.mu.Lock()
	defer c.mu.Unlock()

	// Increment index and assign it to message.
	c.index++
	m.Index = c.index

	// Append message to the topic.
	c.messagesByTopicID[m.TopicID] = append(c.messagesByTopicID[m.TopicID], m)

	// Send to each connection for the topic.
	for _, conn := range c.conns {
		if conn.topicID == m.TopicID {
			conn.Send(m)
		}
	}

	return m.Index, nil
}

func (c *MessagingClient) Conn(topicID uint64) influxdb.MessagingConn {
	return c.ConnFunc(topicID)
}

func (c *MessagingClient) CloseConn(topicID uint64) error {
	c.mu.Lock()
	defer c.mu.Unlock()

	conns := []*MessagingConn{}
	for _, conn := range c.conns {
		if conn.topicID == topicID {
			if err := conn.Close(); err != nil {
				return err
			}
			continue
		}
		conns = append(conns, conn)
	}
	c.conns = conns
	return nil
}

// DefaultConnFunc returns a connection for a specific topic.
func (c *MessagingClient) DefaultConnFunc(topicID uint64) influxdb.MessagingConn {
	c.mu.Lock()
	defer c.mu.Unlock()

	// Create new connection.
	conn := NewMessagingConn(topicID, c.dataURL)

	// Track connections.
	c.conns = append(c.conns, conn)

	return conn
}

// Sync blocks until a given index has been sent through the client.
func (c *MessagingClient) Sync(index uint64) {
	for {
		c.mu.Lock()
		if c.index >= index {
			c.mu.Unlock()
			time.Sleep(10 * time.Millisecond)
			return
		}
		c.mu.Unlock()

		// Otherwise wait momentarily and check again.
		time.Sleep(1 * time.Millisecond)
	}
}

// MessagingConn represents a mockable connection implementing influxdb.MessagingConn.
type MessagingConn struct {
	mu      sync.Mutex
	topicID uint64
	index   uint64
	dataURL url.URL
	c       chan *messaging.Message
}

// NewMessagingConn returns a new instance of MessagingConn.
func NewMessagingConn(topicID uint64, dataURL url.URL) *MessagingConn {
	return &MessagingConn{
		topicID: topicID,
		dataURL: dataURL,
	}
}

// Open starts the stream from a given index.
func (c *MessagingConn) Open(index uint64, streaming bool) error {
	// TODO: Fill connection stream with existing messages.
	c.c = make(chan *messaging.Message, 1024)
	return nil
}

// Close closes the streaming channel.
func (c *MessagingConn) Close() error {
	close(c.c)
	return nil
}

// C returns a channel for streaming message.
func (c *MessagingConn) C() <-chan *messaging.Message { return c.c }

// SetIndex sets the most-recently replicated index on the connection.
func (c *MessagingConn) SetIndex(index uint64) {
}

func (c *MessagingConn) Send(m *messaging.Message) {
	// Ignore any old messages.
	c.mu.Lock()
	if m.Index <= c.index {
		c.mu.Unlock()
		return
	}
	c.index = m.Index
	c.mu.Unlock()

	// Send message to channel.
	c.c <- m
}