Register the broker wait before the first poll (closes #157)

GET /api/v1/messages read the client output queue and only then
registered with the broker. Notify wakes only the channels registered
at that moment, so a message queued between the empty read and the
registration waited out the whole long-poll timeout. The handler now
registers first and removes its channel on every exit path; the IRC
relay loop registers before each drain in the same way.

The handlers read the queue through a function field that defaults to
Database.PollMessages, so tests can act just after a read: one checks
that a message queued there comes back at once, others that no wait
channel stays registered however a request ends.

Model: opus-5-5

internal/handlers/api.go

@@ -813,7 +813,14 @@ func (hdlr *Handlers) HandleGetMessages() http.HandlerFunc {
 			timeout = maxLongPollTimeout
 		}
 
-		msgs, lastQID, err := hdlr.params.Database.PollMessages(
+		// Register with the broker before reading the queue.
+		// Notify wakes only the channels registered at that
+		// moment, so a message queued after the read below
+		// must find this one there already.
+		waitCh := hdlr.broker.Wait(sessionID)
+		defer hdlr.broker.Remove(sessionID, waitCh)
+
+		msgs, lastQID, err := hdlr.pollMessages(
 			request.Context(), clientID,
 			afterID, pollMessageLimit,
 		)
@@ -841,19 +848,22 @@ func (hdlr *Handlers) HandleGetMessages() http.HandlerFunc {
 
 		hdlr.longPoll(
 			writer, request,
-			sessionID, clientID, afterID, timeout,
+			waitCh, clientID, afterID, timeout,
 		)
 	}
 }
 
+// longPoll waits up to timeout seconds for the broker to
+// signal waitCh, then answers with the messages queued for
+// clientID after afterID. It answers nothing if the client
+// goes away first.
 func (hdlr *Handlers) longPoll(
 	writer http.ResponseWriter,
 	request *http.Request,
-	sessionID, clientID, afterID int64,
+	waitCh <-chan struct{},
+	clientID, afterID int64,
 	timeout int,
 ) {
-	waitCh := hdlr.broker.Wait(sessionID)
-
 	timer := time.NewTimer(
 		time.Duration(timeout) * time.Second,
 	)
@@ -864,14 +874,10 @@ func (hdlr *Handlers) longPoll(
 	case <-waitCh:
 	case <-timer.C:
 	case <-request.Context().Done():
-		hdlr.broker.Remove(sessionID, waitCh)
-
 		return
 	}
 
-	hdlr.broker.Remove(sessionID, waitCh)
-
-	msgs, lastQID, err := hdlr.params.Database.PollMessages(
+	msgs, lastQID, err := hdlr.pollMessages(
 		request.Context(), clientID,
 		afterID, pollMessageLimit,
 	)

internal/handlers/api_test.go

@@ -140,6 +140,7 @@ type testServer struct {
 	fxApp      *fxtest.App
 	broker     *broker.Broker
 	database   *db.Database
+	handlers   *handlers.Handlers
 }
 
 func newTestServer(
@@ -170,6 +171,7 @@ func newTestServerWith(
 		srv      *server.Server
 		brk      *broker.Broker
 		database *db.Database
+		hdlr     *handlers.Handlers
 	)
 
 	app := fxtest.New(t,
@@ -210,7 +212,7 @@ func newTestServerWith(
 			newTestHandlers,
 			newTestServerFx,
 		),
-		fx.Populate(&srv, &brk, &database),
+		fx.Populate(&srv, &brk, &database, &hdlr),
 	)
 
 	app.RequireStart()
@@ -228,6 +230,7 @@ func newTestServerWith(
 		fxApp:      app,
 		broker:     brk,
 		database:   database,
+		handlers:   hdlr,
 	}
 }
 
@@ -2554,6 +2557,364 @@ func TestLongPollTimeout(t *testing.T) {
 	}
 }
 
+// TestLongPollReturnsMessageQueuedAfterFirstRead checks that
+// a message queued, and the broker notified, after the long
+// poll's first read of the queue found nothing but before the
+// poll waits, is returned at once and not when the poll times
+// out, and that the poll then holds no broker wait channel.
+func TestLongPollReturnsMessageQueuedAfterFirstRead(t *testing.T) {
+	const (
+		nick   = "lp_late"
+		sender = "lp_late_sender"
+
+		// pollTimeout is the long poll's timeout, in seconds.
+		pollTimeout = 10
+
+		// atOnce is the longest the long poll may take to
+		// return the message.
+		atOnce = 5 * time.Second
+	)
+
+	tserver := newTestServer(t)
+	token := tserver.createSession(nick)
+
+	_, lastID := tserver.pollMessages(token, 0)
+
+	sessionID, err := tserver.database.GetSessionByNick(t.Context(), nick)
+	if err != nil {
+		t.Fatalf("get session of %s: %v", nick, err)
+	}
+
+	queued := make(chan error, 1)
+
+	var once sync.Once
+
+	tserver.handlers.SetPollMessages(func(
+		ctx context.Context,
+		clientID, afterQueueID int64,
+		limit int,
+	) ([]db.IRCMessage, int64, error) {
+		msgs, lastQID, pollErr := tserver.database.PollMessages(
+			ctx, clientID, afterQueueID, limit,
+		)
+
+		once.Do(func() {
+			queued <- tserver.queueMessage(ctx, sessionID, sender, nick)
+		})
+
+		return msgs, lastQID, pollErr
+	})
+
+	start := time.Now()
+
+	msgs, err := tserver.longPoll(token, lastID, pollTimeout)
+	if err != nil {
+		t.Fatalf("long poll: %v", err)
+	}
+
+	elapsed := time.Since(start)
+
+	select {
+	case queueErr := <-queued:
+		if queueErr != nil {
+			t.Fatalf("queue message: %v", queueErr)
+		}
+	default:
+		t.Fatal("the long poll did not read the queue")
+	}
+
+	if elapsed > atOnce {
+		t.Fatalf("the long poll returned after %v, want within %v",
+			elapsed, atOnce)
+	}
+
+	if !findMessage(msgs, privmsgCmd, sender) {
+		t.Fatalf("long poll: got %v, want the PRIVMSG from %s",
+			msgs, sender)
+	}
+
+	tserver.checkWaitChannelRemoved(sessionID)
+}
+
+// errReadFailed is what a test's replacement for
+// Database.PollMessages returns to make a read of the queue
+// fail.
+var errReadFailed = errors.New("read failed")
+
+// TestGetMessagesRemovesWaitChannel checks that a request to
+// GET /api/v1/messages leaves no broker wait channel behind,
+// whether it is answered at once, after its timeout, or with
+// an error because a read of the queue failed.
+func TestGetMessagesRemovesWaitChannel(t *testing.T) {
+	const (
+		// longTimeout, in seconds, is the timeout of the
+		// requests that must be answered without waiting.
+		longTimeout = 10
+
+		// shortTimeout, in seconds, is the timeout of the
+		// requests that wait it out.
+		shortTimeout = 1
+	)
+
+	// queued is whether a message is queued before the request,
+	// whose timeout is in seconds. failingRead is as for
+	// getMessagesFailingRead. wantReads is how many times the
+	// handler reads the queue: once if it answers at once,
+	// twice if it waits first.
+	tests := []struct {
+		name                  string
+		queued                bool
+		timeout, failingRead  int
+		wantReads, wantStatus int
+	}{
+		{
+			name: "nothing queued, timeout 0", timeout: 0,
+			wantReads: 1, wantStatus: http.StatusOK,
+		},
+		{
+			name: "message queued", queued: true, timeout: longTimeout,
+			wantReads: 1, wantStatus: http.StatusOK,
+		},
+		{
+			name: "timed out", timeout: shortTimeout,
+			wantReads: 2, wantStatus: http.StatusOK,
+		},
+		{
+			name: "first read failed", timeout: longTimeout, failingRead: 1,
+			wantReads: 1, wantStatus: http.StatusInternalServerError,
+		},
+		{
+			name: "read after the wait failed", timeout: shortTimeout,
+			failingRead: 2, wantReads: 2,
+			wantStatus: http.StatusInternalServerError,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			const nick = "wc_poller"
+
+			tserver := newTestServer(t)
+			token := tserver.createSession(nick)
+
+			_, lastID := tserver.pollMessages(token, 0)
+
+			sessionID, err := tserver.database.GetSessionByNick(
+				t.Context(), nick,
+			)
+			if err != nil {
+				t.Fatalf("get session of %s: %v", nick, err)
+			}
+
+			if tc.queued {
+				err = tserver.queueMessage(
+					t.Context(), sessionID, "wc_sender", nick,
+				)
+				if err != nil {
+					t.Fatalf("queue message: %v", err)
+				}
+			}
+
+			status, reads := tserver.getMessagesFailingRead(
+				token, lastID, tc.timeout, tc.failingRead,
+			)
+			if status != tc.wantStatus {
+				t.Errorf("status %d, want %d", status, tc.wantStatus)
+			}
+
+			if reads != tc.wantReads {
+				t.Errorf("the handler read the queue %d times, want %d",
+					reads, tc.wantReads)
+			}
+
+			tserver.checkWaitChannelRemoved(sessionID)
+		})
+	}
+}
+
+// getMessagesFailingRead sends GET /api/v1/messages for the
+// client token with timeout and after afterID. The handler's
+// read of the queue numbered failingRead, counting from 1,
+// fails; with 0, none does. It returns the response status
+// and how many times the handler read the queue.
+func (tserver *testServer) getMessagesFailingRead(
+	token string,
+	afterID int64,
+	timeout, failingRead int,
+) (int, int) {
+	tserver.t.Helper()
+
+	reads := 0
+
+	tserver.handlers.SetPollMessages(func(
+		ctx context.Context,
+		clientID, afterQueueID int64,
+		limit int,
+	) ([]db.IRCMessage, int64, error) {
+		reads++
+
+		if reads == failingRead {
+			return nil, 0, errReadFailed
+		}
+
+		return tserver.database.PollMessages(
+			ctx, clientID, afterQueueID, limit,
+		)
+	})
+
+	resp, err := doRequestAuth(
+		tserver.t, http.MethodGet,
+		tserver.url(fmt.Sprintf(
+			"%s?timeout=%d&after=%d", apiMessages, timeout, afterID,
+		)),
+		token, nil,
+	)
+	if err != nil {
+		tserver.t.Fatalf("get messages: %v", err)
+	}
+
+	_ = resp.Body.Close()
+
+	return resp.StatusCode, reads
+}
+
+// TestLongPollRemovesWaitChannelWhenClientGivesUp checks that
+// a long poll whose client gives up while it waits leaves no
+// broker wait channel behind.
+func TestLongPollRemovesWaitChannelWhenClientGivesUp(t *testing.T) {
+	const (
+		nick = "wc_leaver"
+
+		// pollTimeout is the long poll's timeout, in seconds.
+		pollTimeout = 10
+
+		// readDeadline is how long the handler has to read
+		// the queue.
+		readDeadline = 5 * time.Second
+	)
+
+	tserver := newTestServer(t)
+	token := tserver.createSession(nick)
+
+	_, lastID := tserver.pollMessages(token, 0)
+
+	sessionID, err := tserver.database.GetSessionByNick(t.Context(), nick)
+	if err != nil {
+		t.Fatalf("get session of %s: %v", nick, err)
+	}
+
+	read := make(chan struct{})
+
+	var once sync.Once
+
+	tserver.handlers.SetPollMessages(func(
+		ctx context.Context,
+		clientID, afterQueueID int64,
+		limit int,
+	) ([]db.IRCMessage, int64, error) {
+		msgs, lastQID, pollErr := tserver.database.PollMessages(
+			ctx, clientID, afterQueueID, limit,
+		)
+
+		once.Do(func() { close(read) })
+
+		return msgs, lastQID, pollErr
+	})
+
+	ctx, giveUp := context.WithCancel(t.Context())
+	defer giveUp()
+
+	request, err := http.NewRequestWithContext(
+		ctx, http.MethodGet,
+		tserver.url(fmt.Sprintf(
+			"%s?timeout=%d&after=%d", apiMessages, pollTimeout, lastID,
+		)),
+		nil,
+	)
+	if err != nil {
+		t.Fatalf("new request: %v", err)
+	}
+
+	request.AddCookie(&http.Cookie{ //nolint:exhaustruct,gosec // only name+value
+		Name:  authCookieName,
+		Value: token,
+	})
+
+	done := make(chan struct{})
+
+	go func() {
+		defer close(done)
+
+		resp, doErr := http.DefaultClient.Do(request)
+		if doErr == nil {
+			_ = resp.Body.Close()
+		}
+	}()
+
+	select {
+	case <-read:
+	case <-time.After(readDeadline):
+		t.Fatal("the long poll did not read the queue")
+	}
+
+	giveUp()
+	<-done
+
+	tserver.checkWaitChannelRemoved(sessionID)
+}
+
+// checkWaitChannelRemoved fails the test unless the broker
+// holds no wait channel for the session sessionID. A client
+// can have its response before the handler has returned, and
+// a handler notices a client that gives up only later, so it
+// waits up to removeDeadline for that.
+func (tserver *testServer) checkWaitChannelRemoved(sessionID int64) {
+	tserver.t.Helper()
+
+	const (
+		removeDeadline = 5 * time.Second
+		checkStep      = 10 * time.Millisecond
+	)
+
+	deadline := time.Now().Add(removeDeadline)
+
+	for tserver.broker.ListenerCount(sessionID) != 0 {
+		if time.Now().After(deadline) {
+			tserver.t.Fatalf("got %d broker wait channels, want 0",
+				tserver.broker.ListenerCount(sessionID))
+		}
+
+		time.Sleep(checkStep)
+	}
+}
+
+// queueMessage stores a PRIVMSG from sender to nick, queues
+// it to the session sessionID and notifies the broker, as the
+// server does for a DM. It returns errors instead of failing
+// the test, so a handler's goroutine can call it.
+func (tserver *testServer) queueMessage(
+	ctx context.Context,
+	sessionID int64,
+	sender, nick string,
+) error {
+	msgID, _, err := tserver.database.InsertMessage(
+		ctx, privmsgCmd, sender, nick, sessionID, 0,
+		nil, json.RawMessage(`["hello"]`), nil,
+	)
+	if err != nil {
+		return fmt.Errorf("insert message: %w", err)
+	}
+
+	err = tserver.database.EnqueueToSession(ctx, sessionID, msgID)
+	if err != nil {
+		return fmt.Errorf("enqueue: %w", err)
+	}
+
+	tserver.broker.Notify(sessionID)
+
+	return nil
+}
+
 func TestEphemeralChannelCleanup(t *testing.T) {
 	const channel = "#ephemeral"
 

internal/handlers/export_test.go

@@ -0,0 +1,19 @@
+package handlers
+
+import (
+	"context"
+
+	"sneak.berlin/go/neoirc/internal/db"
+)
+
+// SetPollMessages makes HandleGetMessages read the client
+// output queue with poll in place of Database.PollMessages.
+func (hdlr *Handlers) SetPollMessages(
+	poll func(
+		ctx context.Context,
+		clientID, afterQueueID int64,
+		limit int,
+	) ([]db.IRCMessage, int64, error),
+) {
+	hdlr.pollMessages = poll
+}

internal/handlers/handlers.go

@@ -64,6 +64,16 @@ type Handlers struct {
 	loginLimiter  *ratelimit.Limiter
 	stats         *stats.Tracker
 	cancelCleanup context.CancelFunc
+
+	// pollMessages is Database.PollMessages, with which
+	// HandleGetMessages reads the client output queue. Tests
+	// replace it to act just after a read, or to make a read
+	// fail.
+	pollMessages func(
+		ctx context.Context,
+		clientID, afterQueueID int64,
+		limit int,
+	) ([]db.IRCMessage, int64, error)
 }
 
 // New creates a new Handlers instance.
@@ -95,6 +105,7 @@ func New(
 		hashcashVal:  hashcash.NewValidator(resource),
 		loginLimiter: ratelimit.New(loginRate, loginBurst),
 		stats:        params.Stats,
+		pollMessages: params.Database.PollMessages,
 	}
 
 	lifecycle.Append(fx.Hook{

internal/ircserver/relay.go

@@ -27,17 +27,22 @@ func (c *Conn) relayMessages(ctx context.Context) {
 		default:
 		}
 
+		// Register with the broker before draining, so that a
+		// message queued after the drain has read the queue
+		// still signals waitCh.
+		waitCh := c.brk.Wait(c.sessionID)
+
 		// Drain any available messages.
 		delivered := c.drainQueue(ctx)
 
 		if delivered {
+			c.brk.Remove(c.sessionID, waitCh)
+
 			// Tight loop while there are messages.
 			continue
 		}
 
 		// Wait for notification or timeout.
-		waitCh := c.brk.Wait(c.sessionID)
-
 		select {
 		case <-waitCh:
 			// New message notification — loop back.

internal/ircserver/relay_test.go

@@ -47,14 +47,13 @@ func (env *testEnv) queueMessage(
 	env.brk.Notify(sessionID)
 }
 
-// TestRelayKeepsOneWaitChannel checks that an idle IRC
-// listener connection holds at most one broker wait channel
-// across several relay ticks, and none once it closes.
+// TestRelayKeepsOneWaitChannel checks that an IRC listener
+// connection holds at most one broker wait channel across
+// several relay ticks, both while idle and after the relay
+// has delivered it a message, and none once it closes.
 func TestRelayKeepsOneWaitChannel(t *testing.T) {
 	t.Parallel()
 
-	const idleTicks = 5
-
 	env := newTestEnv(t)
 
 	client := env.dial(t)
@@ -65,15 +64,17 @@ func TestRelayKeepsOneWaitChannel(t *testing.T) {
 		t.Fatalf("get session of idle: %v", err)
 	}
 
-	for tick := range idleTicks {
-		time.Sleep(ircserver.PollInterval)
+	env.checkOneWaitChannel(t, sessionID)
 
-		got := env.brk.ListenerCount(sessionID)
-		if got > 1 {
-			t.Fatalf("tick %d: got %d wait channels, want at most 1",
-				tick+1, got)
-		}
-	}
+	env.queueMessage(
+		t, "idle", "PRIVMSG", "carol", "idle", []string{"hello"},
+	)
+
+	client.readUntil(func(l string) bool {
+		return strings.Contains(l, "hello")
+	})
+
+	env.checkOneWaitChannel(t, sessionID)
 
 	err = client.conn.Close()
 	if err != nil {
@@ -92,6 +93,28 @@ func TestRelayKeepsOneWaitChannel(t *testing.T) {
 	}
 }
 
+// checkOneWaitChannel fails the test if the session
+// sessionID holds more than one broker wait channel at any of
+// several relay ticks.
+func (env *testEnv) checkOneWaitChannel(
+	t *testing.T,
+	sessionID int64,
+) {
+	t.Helper()
+
+	const ticks = 5
+
+	for tick := range ticks {
+		time.Sleep(ircserver.PollInterval)
+
+		got := env.brk.ListenerCount(sessionID)
+		if got > 1 {
+			t.Fatalf("tick %d: got %d wait channels, want at most 1",
+				tick+1, got)
+		}
+	}
+}
+
 // TestRelayMultiLineBody checks that each line of a PRIVMSG
 // or NOTICE body sent over the HTTP API reaches an IRC
 // listener user as its own line.