open-consul/agent/consul/stream/event_publisher_test.go

package stream

import (
	"context"
	"fmt"
	"testing"
	"time"

	"github.com/stretchr/testify/require"

	"github.com/hashicorp/consul/acl"
	"github.com/hashicorp/consul/proto/private/pbsubscribe"
	"github.com/hashicorp/consul/sdk/testutil"
)

type intTopic int

func (i intTopic) String() string {
	return fmt.Sprintf("%d", i)
}

var testTopic Topic = intTopic(999)

func TestEventPublisher_SubscribeWithIndex0(t *testing.T) {
	req := &SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	publisher := NewEventPublisher(0)
	registerTestSnapshotHandlers(t, publisher)
	go publisher.Run(ctx)

	sub, err := publisher.Subscribe(req)
	require.NoError(t, err)
	defer sub.Unsubscribe()
	eventCh := runSubscription(ctx, sub)

	next := getNextEvent(t, eventCh)
	require.Equal(t, testSnapshotEvent, next)

	next = getNextEvent(t, eventCh)
	require.True(t, next.IsEndOfSnapshot())

	assertNoResult(t, eventCh)

	events := []Event{{
		Topic:   testTopic,
		Payload: simplePayload{key: "sub-key", value: "the-published-event-payload"},
	}}
	publisher.Publish(events)

	// Subscriber should see the published event
	next = getNextEvent(t, eventCh)
	expected := Event{
		Topic:   testTopic,
		Payload: simplePayload{key: "sub-key", value: "the-published-event-payload"},
	}
	require.Equal(t, expected, next)

	// Subscriber should not see events for other keys
	publisher.Publish([]Event{{
		Topic:   testTopic,
		Payload: simplePayload{key: "other-key", value: "this-should-not-reach-the-subscriber"},
	}})
	assertNoResult(t, eventCh)
}

var testSnapshotEvent = Event{
	Topic:   testTopic,
	Payload: simplePayload{key: "sub-key", value: "snapshot-event-payload"},
	Index:   1,
}

type simplePayload struct {
	key        string
	value      string
	noReadPerm bool
}

func (p simplePayload) HasReadPermission(acl.Authorizer) bool {
	return !p.noReadPerm
}

func (p simplePayload) Subject() Subject { return StringSubject(p.key) }

func (p simplePayload) ToSubscriptionEvent(idx uint64) *pbsubscribe.Event {
	panic("simplePayload does not implement ToSubscriptionEvent")
}

func registerTestSnapshotHandlers(t *testing.T, publisher *EventPublisher) {
	t.Helper()

	testTopicHandler := func(req SubscribeRequest, buf SnapshotAppender) (uint64, error) {
		if req.Topic != testTopic {
			return 0, fmt.Errorf("unexpected topic: %v", req.Topic)
		}
		buf.Append([]Event{testSnapshotEvent})
		return 1, nil
	}

	require.NoError(t, publisher.RegisterHandler(testTopic, testTopicHandler, false))
}

func runSubscription(ctx context.Context, sub *Subscription) <-chan eventOrErr {
	eventCh := make(chan eventOrErr, 1)
	go func() {
		for {
			es, err := sub.Next(ctx)
			eventCh <- eventOrErr{
				Event: es,
				Err:   err,
			}
			if err != nil {
				return
			}
		}
	}()
	return eventCh
}

type eventOrErr struct {
	Event Event
	Err   error
}

func getNextEvent(t *testing.T, eventCh <-chan eventOrErr) Event {
	t.Helper()
	select {
	case next := <-eventCh:
		require.NoError(t, next.Err)
		return next.Event
	case <-time.After(100 * time.Millisecond):
		t.Fatalf("timeout waiting for event from subscription")
		return Event{}
	}
}

func assertNoResult(t *testing.T, eventCh <-chan eventOrErr) {
	t.Helper()
	select {
	case next := <-eventCh:
		require.NoError(t, next.Err)
		t.Fatalf("received unexpected event: %#v", next.Event.Payload)
	case <-time.After(25 * time.Millisecond):
	}
}

func TestEventPublisher_ShutdownClosesSubscriptions(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	fn := func(req SubscribeRequest, buf SnapshotAppender) (uint64, error) {
		return 0, nil
	}

	publisher := NewEventPublisher(time.Second)
	registerTestSnapshotHandlers(t, publisher)
	publisher.RegisterHandler(intTopic(22), fn, false)
	publisher.RegisterHandler(intTopic(33), fn, false)
	go publisher.Run(ctx)

	sub1, err := publisher.Subscribe(&SubscribeRequest{Topic: intTopic(22), Subject: SubjectNone})
	require.NoError(t, err)
	defer sub1.Unsubscribe()

	sub2, err := publisher.Subscribe(&SubscribeRequest{Topic: intTopic(33), Subject: SubjectNone})
	require.NoError(t, err)
	defer sub2.Unsubscribe()

	cancel() // Shutdown

	err = consumeSub(context.Background(), sub1)
	require.Equal(t, err, ErrShuttingDown)

	_, err = sub2.Next(context.Background())
	require.Equal(t, err, ErrShuttingDown)
}

func consumeSub(ctx context.Context, sub *Subscription) error {
	for {
		event, err := sub.Next(ctx)
		switch {
		case err != nil:
			return err
		case event.IsEndOfSnapshot():
			continue
		}
	}
}

func TestEventPublisher_SubscribeWithIndex0_FromCache(t *testing.T) {
	req := &SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	publisher := NewEventPublisher(time.Second)
	registerTestSnapshotHandlers(t, publisher)
	go publisher.Run(ctx)

	sub, err := publisher.Subscribe(req)
	require.NoError(t, err)
	defer sub.Unsubscribe()

	publisher.snapshotHandlers[testTopic] = func(_ SubscribeRequest, _ SnapshotAppender) (uint64, error) {
		return 0, fmt.Errorf("error should not be seen, cache should have been used")
	}

	sub, err = publisher.Subscribe(req)
	require.NoError(t, err)
	defer sub.Unsubscribe()

	eventCh := runSubscription(ctx, sub)
	next := getNextEvent(t, eventCh)
	require.Equal(t, testSnapshotEvent, next)

	next = getNextEvent(t, eventCh)
	require.True(t, next.IsEndOfSnapshot())

	// Now subscriber should block waiting for updates
	assertNoResult(t, eventCh)

	expected := Event{
		Topic:   testTopic,
		Payload: simplePayload{key: "sub-key", value: "the-published-event-payload"},
		Index:   3,
	}
	publisher.Publish([]Event{expected})

	// Subscriber should see the published event
	next = getNextEvent(t, eventCh)
	require.Equal(t, expected, next)
}

func TestEventPublisher_SubscribeWithIndexNotZero_CanResume(t *testing.T) {
	req := &SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	publisher := NewEventPublisher(time.Second)
	registerTestSnapshotHandlers(t, publisher)
	go publisher.Run(ctx)

	simulateExistingSubscriber(t, publisher, req)

	// Publish the testSnapshotEvent, to ensure that it is skipped over when
	// splicing the topic buffer onto the snapshot.
	publisher.publishEvent([]Event{testSnapshotEvent})

	testutil.RunStep(t, "start a subscription and unsub", func(t *testing.T) {
		sub, err := publisher.Subscribe(req)
		require.NoError(t, err)
		defer sub.Unsubscribe()

		eventCh := runSubscription(ctx, sub)

		next := getNextEvent(t, eventCh)
		require.Equal(t, testSnapshotEvent, next)

		next = getNextEvent(t, eventCh)
		require.True(t, next.IsEndOfSnapshot())
		require.Equal(t, uint64(1), next.Index)
	})

	testutil.RunStep(t, "resume the subscription", func(t *testing.T) {
		newReq := *req
		newReq.Index = 1
		sub, err := publisher.Subscribe(&newReq)
		require.NoError(t, err)

		eventCh := runSubscription(ctx, sub)
		assertNoResult(t, eventCh)

		expected := Event{
			Topic:   testTopic,
			Index:   3,
			Payload: simplePayload{key: "sub-key", value: "event-3"},
		}
		publisher.publishEvent([]Event{expected})

		next := getNextEvent(t, eventCh)
		require.Equal(t, expected, next)
	})
}

func TestEventPublisher_SubscribeWithIndexNotZero_NewSnapshot(t *testing.T) {
	req := &SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	publisher := NewEventPublisher(0)
	registerTestSnapshotHandlers(t, publisher)
	go publisher.Run(ctx)
	// Include the same event in the topicBuffer
	publisher.publishEvent([]Event{testSnapshotEvent})

	testutil.RunStep(t, "start a subscription and unsub", func(t *testing.T) {
		sub, err := publisher.Subscribe(req)
		require.NoError(t, err)
		defer sub.Unsubscribe()

		eventCh := runSubscription(ctx, sub)

		next := getNextEvent(t, eventCh)
		require.Equal(t, testSnapshotEvent, next)

		next = getNextEvent(t, eventCh)
		require.True(t, next.IsEndOfSnapshot())
		require.Equal(t, uint64(1), next.Index)
	})

	nextEvent := Event{
		Topic:   testTopic,
		Index:   3,
		Payload: simplePayload{key: "sub-key", value: "event-3"},
	}

	testutil.RunStep(t, "publish an event while unsubed", func(t *testing.T) {
		publisher.publishEvent([]Event{nextEvent})
	})

	testutil.RunStep(t, "resume the subscription", func(t *testing.T) {
		newReq := *req
		newReq.Index = 1
		sub, err := publisher.Subscribe(&newReq)
		require.NoError(t, err)

		eventCh := runSubscription(ctx, sub)
		next := getNextEvent(t, eventCh)
		require.True(t, next.IsNewSnapshotToFollow(), next)

		next = getNextEvent(t, eventCh)
		require.Equal(t, testSnapshotEvent, next)

		next = getNextEvent(t, eventCh)
		require.True(t, next.IsEndOfSnapshot())
	})
}

func TestEventPublisher_SubscribeWithIndexNotZero_NewSnapshotFromCache(t *testing.T) {
	req := &SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
	}
	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	publisher := NewEventPublisher(time.Second)
	registerTestSnapshotHandlers(t, publisher)
	go publisher.Run(ctx)

	simulateExistingSubscriber(t, publisher, req)

	// Publish the testSnapshotEvent, to ensure that it is skipped over when
	// splicing the topic buffer onto the snapshot.
	publisher.publishEvent([]Event{testSnapshotEvent})

	testutil.RunStep(t, "start a subscription and unsub", func(t *testing.T) {
		sub, err := publisher.Subscribe(req)
		require.NoError(t, err)
		defer sub.Unsubscribe()

		eventCh := runSubscription(ctx, sub)

		next := getNextEvent(t, eventCh)
		require.Equal(t, testSnapshotEvent, next)

		next = getNextEvent(t, eventCh)
		require.True(t, next.IsEndOfSnapshot())
		require.Equal(t, uint64(1), next.Index)
	})

	nextEvent := Event{
		Topic:   testTopic,
		Index:   3,
		Payload: simplePayload{key: "sub-key", value: "event-3"},
	}

	testutil.RunStep(t, "publish an event while unsubed", func(t *testing.T) {
		publisher.publishEvent([]Event{nextEvent})
	})

	publisher.snapshotHandlers[testTopic] = func(_ SubscribeRequest, _ SnapshotAppender) (uint64, error) {
		return 0, fmt.Errorf("error should not be seen, cache should have been used")
	}

	testutil.RunStep(t, "resume the subscription", func(t *testing.T) {
		newReq := *req
		newReq.Index = 1
		sub, err := publisher.Subscribe(&newReq)
		require.NoError(t, err)
		defer sub.Unsubscribe()

		eventCh := runSubscription(ctx, sub)
		next := getNextEvent(t, eventCh)
		require.True(t, next.IsNewSnapshotToFollow(), next)

		next = getNextEvent(t, eventCh)
		require.Equal(t, testSnapshotEvent, next)

		next = getNextEvent(t, eventCh)
		require.True(t, next.IsEndOfSnapshot())

		next = getNextEvent(t, eventCh)
		require.Equal(t, nextEvent, next)
	})
}

func TestEventPublisher_SubscribeWithIndexNotZero_NewSnapshot_WithCache(t *testing.T) {
	req := &SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
		Index:   1,
	}

	nextEvent := Event{
		Topic:   testTopic,
		Index:   3,
		Payload: simplePayload{key: "sub-key", value: "event-3"},
	}

	testTopicHandler := func(req SubscribeRequest, buf SnapshotAppender) (uint64, error) {
		if req.Topic != testTopic {
			return 0, fmt.Errorf("unexpected topic: %v", req.Topic)
		}
		buf.Append([]Event{testSnapshotEvent})
		buf.Append([]Event{nextEvent})
		return 3, nil
	}

	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()

	publisher := NewEventPublisher(time.Second)
	publisher.RegisterHandler(testTopic, testTopicHandler, false)
	go publisher.Run(ctx)

	simulateExistingSubscriber(t, publisher, req)

	// Publish the events, to ensure they are is skipped over when splicing the
	// topic buffer onto the snapshot.
	publisher.publishEvent([]Event{testSnapshotEvent})
	publisher.publishEvent([]Event{nextEvent})

	testutil.RunStep(t, "start a subscription and unsub", func(t *testing.T) {
		sub, err := publisher.Subscribe(req)
		require.NoError(t, err)
		defer sub.Unsubscribe()

		eventCh := runSubscription(ctx, sub)
		next := getNextEvent(t, eventCh)
		require.True(t, next.IsNewSnapshotToFollow(), next)

		next = getNextEvent(t, eventCh)
		require.Equal(t, testSnapshotEvent, next)

		next = getNextEvent(t, eventCh)
		require.Equal(t, nextEvent, next)

		next = getNextEvent(t, eventCh)
		require.True(t, next.IsEndOfSnapshot(), next)
		require.Equal(t, uint64(3), next.Index)
	})

	publisher.snapshotHandlers[testTopic] = func(_ SubscribeRequest, _ SnapshotAppender) (uint64, error) {
		return 0, fmt.Errorf("error should not be seen, cache should have been used")
	}

	testutil.RunStep(t, "resume the subscription", func(t *testing.T) {
		newReq := *req
		newReq.Index = 0
		sub, err := publisher.Subscribe(&newReq)
		require.NoError(t, err)

		eventCh := runSubscription(ctx, sub)
		next := getNextEvent(t, eventCh)
		require.Equal(t, testSnapshotEvent, next)

		next = getNextEvent(t, eventCh)
		require.Equal(t, nextEvent, next)

		next = getNextEvent(t, eventCh)
		require.True(t, next.IsEndOfSnapshot())
	})
}

func TestEventPublisher_Unsubscribe_ClosesSubscription(t *testing.T) {
	req := &SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
	}
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	publisher := NewEventPublisher(time.Second)
	registerTestSnapshotHandlers(t, publisher)

	sub, err := publisher.Subscribe(req)
	require.NoError(t, err)

	_, err = sub.Next(ctx)
	require.NoError(t, err)

	sub.Unsubscribe()
	_, err = sub.Next(ctx)
	require.Error(t, err)
	require.Contains(t, err.Error(), "subscription was closed by unsubscribe")
}

func TestEventPublisher_Unsubscribe_FreesResourcesWhenThereAreNoSubscribers(t *testing.T) {
	req := &SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
	}

	publisher := NewEventPublisher(time.Second)
	registerTestSnapshotHandlers(t, publisher)

	sub1, err := publisher.Subscribe(req)
	require.NoError(t, err)

	// Expect a topic buffer and snapshot to have been created.
	publisher.lock.Lock()
	require.NotNil(t, publisher.topicBuffers[req.topicSubject()])
	require.NotNil(t, publisher.snapCache[req.topicSubject()])
	publisher.lock.Unlock()

	// Create another subscription and close the old one, to ensure the buffer and
	// snapshot stick around as long as there's at least one subscriber.
	sub2, err := publisher.Subscribe(req)
	require.NoError(t, err)

	sub1.Unsubscribe()

	publisher.lock.Lock()
	require.NotNil(t, publisher.topicBuffers[req.topicSubject()])
	require.NotNil(t, publisher.snapCache[req.topicSubject()])
	publisher.lock.Unlock()

	// Close the other subscription and expect the buffer and snapshot to have
	// been cleaned up.
	sub2.Unsubscribe()

	publisher.lock.Lock()
	require.Nil(t, publisher.topicBuffers[req.topicSubject()])
	require.Nil(t, publisher.snapCache[req.topicSubject()])
	publisher.lock.Unlock()
}

// simulateExistingSubscriber creates a subscription that remains open throughout
// a test to prevent the topic buffer getting garbage-collected.
//
// It evicts the created snapshot from the cache immediately (simulating an
// existing subscription that has been open long enough the snapshot's TTL has
// been reached) so you can test snapshots getting created afresh.
func simulateExistingSubscriber(t *testing.T, p *EventPublisher, r *SubscribeRequest) {
	t.Helper()

	sub, err := p.Subscribe(r)
	require.NoError(t, err)
	t.Cleanup(sub.Unsubscribe)

	p.lock.Lock()
	delete(p.snapCache, r.topicSubject())
	p.lock.Unlock()
}

func TestEventPublisher_Subscribe_WildcardNotSupported(t *testing.T) {
	publisher := NewEventPublisher(0)

	handler := func(SubscribeRequest, SnapshotAppender) (uint64, error) { return 0, nil }
	require.NoError(t, publisher.RegisterHandler(testTopic, handler, false))

	_, err := publisher.Subscribe(&SubscribeRequest{
		Topic:   testTopic,
		Subject: SubjectWildcard,
	})
	require.Error(t, err)
	require.Contains(t, err.Error(), "does not support wildcard subscriptions")
}

func TestEventPublisher_Subscribe_WildcardSupported(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	publisher := NewEventPublisher(0)
	go publisher.Run(ctx)

	var (
		// These events are in the snapshot.
		a1 = Event{
			Topic:   testTopic,
			Payload: simplePayload{key: "a", value: "1"},
			Index:   1,
		}
		b1 = Event{
			Topic:   testTopic,
			Payload: simplePayload{key: "b", value: "1"},
			Index:   1,
		}

		// These events are published after the subscription begins.
		a2 = Event{
			Topic:   testTopic,
			Payload: simplePayload{key: "a", value: "2"},
			Index:   2,
		}
		b2 = Event{
			Topic:   testTopic,
			Payload: simplePayload{key: "b", value: "2"},
			Index:   2,
		}
	)

	handler := func(_ SubscribeRequest, buf SnapshotAppender) (uint64, error) {
		buf.Append([]Event{a1, b1})
		return 1, nil
	}
	require.NoError(t, publisher.RegisterHandler(testTopic, handler, true))

	sub, err := publisher.Subscribe(&SubscribeRequest{
		Topic:   testTopic,
		Subject: SubjectWildcard,
	})
	require.NoError(t, err)
	t.Cleanup(sub.Unsubscribe)

	eventCh := runSubscription(ctx, sub)

	next := getNextEvent(t, eventCh)
	require.Equal(t, &PayloadEvents{
		Items: []Event{a1, b1},
	}, next.Payload)

	next = getNextEvent(t, eventCh)
	require.True(t, next.IsEndOfSnapshot(), "expected end of snapshot")

	publisher.Publish([]Event{a2, b2})
	next = getNextEvent(t, eventCh)
	require.Equal(t, &PayloadEvents{
		Items: []Event{a2, b2},
	}, next.Payload)
}

func TestEventPublisher_Publish_WildcardNotAllowed(t *testing.T) {
	publisher := NewEventPublisher(0)

	require.Panics(t, func() {
		publisher.Publish([]Event{
			{
				Topic:   testTopic,
				Payload: wildcardPayload{},
			},
		})
	})
}

type wildcardPayload struct{}

func (wildcardPayload) Subject() Subject                              { return SubjectWildcard }
func (wildcardPayload) HasReadPermission(acl.Authorizer) bool         { return true }
func (wildcardPayload) ToSubscriptionEvent(uint64) *pbsubscribe.Event { return &pbsubscribe.Event{} }

func TestEventPublisher_SnapshotIndex0(t *testing.T) {
	ctx, cancel := context.WithCancel(context.Background())
	t.Cleanup(cancel)

	publisher := NewEventPublisher(10 * time.Second)
	go publisher.Run(ctx)

	publisher.RegisterHandler(testTopic, func(SubscribeRequest, SnapshotAppender) (uint64, error) {
		return 0, nil
	}, false)

	sub, err := publisher.Subscribe(&SubscribeRequest{
		Topic:   testTopic,
		Subject: StringSubject("sub-key"),
	})
	require.NoError(t, err)
	t.Cleanup(sub.Unsubscribe)

	eventCh := runSubscription(ctx, sub)
	event := getNextEvent(t, eventCh)
	require.True(t, event.IsEndOfSnapshot())

	// Even though the snapshot handler returned 0, the subscriber shouldn't see it.
	require.Equal(t, uint64(1), event.Index)
}