diff --git a/internal/lsp/client.go b/internal/lsp/client.go index bf821223c..0c630558a 100644 --- a/internal/lsp/client.go +++ b/internal/lsp/client.go @@ -13,8 +13,14 @@ import ( ) // NotificationHandler receives server->client notifications (e.g. -// textDocument/publishDiagnostics). params is the raw JSON payload. -type NotificationHandler func(method string, params json.RawMessage) +// textDocument/publishDiagnostics). params is the raw JSON payload. seq is the +// notification's receipt sequence (see Client.ReceiptSeq): the position of the +// frame it arrived in, stamped the moment that frame came off the wire — NOT +// when this handler happens to run, which can lag receipt when the queue is +// backed up. A caller that needs to know whether something newer than a given +// point has arrived must compare against seq, not against when its own handling +// code runs. +type NotificationHandler func(method string, params json.RawMessage, seq int64) // Client speaks JSON-RPC 2.0 with LSP framing (Content-Length headers) over a // reader/writer pair. It is transport-agnostic: server.go wires it to a process's @@ -32,8 +38,60 @@ type Client struct { closeOnce sync.Once closed chan struct{} readErr error + + notifyMu sync.Mutex + notifyQueue []notification + notifyBytes int + notifyReady chan struct{} + notifyClosed bool + // notificationDrained is closed by notificationLoop, and only after + // closeNotifications has stopped admission and every accepted notification + // (including an in-flight handler) has finished. + notificationDrained chan struct{} + receiptSeq int64 // frames read off the wire so far (see ReceiptSeq) + // acceptedSeq is the receipt seq of the last notification queued for the + // worker, handledSeq the seq of the last one whose handler has returned, and + // handledWait a channel closed (and replaced) each time handledSeq advances. + // Together they let a waiter tell whether the worker has caught up with what + // the reader accepted. Both track notifications only; receiptSeq counts every + // frame, so the two are compared against each other, never against it. + acceptedSeq int64 + handledSeq int64 + handledWait chan struct{} +} + +type notification struct { + method string + params json.RawMessage + seq int64 } +// The notification backlog is bounded by both message count and retained +// payload bytes. A well-behaved handler drains far faster than any single burst +// fills either limit; sustained overload — a language server emitting faster +// than the single handler can consume, or a handler stuck waiting on a +// re-entrant Call — is a fatal condition for this client, not something to +// paper over by growing the queue without bound. Hitting either limit fails the +// client (see enqueueNotification): IsClosed becomes true, and the manager +// evicts and restarts the session on next use, exactly as it does for any other +// dead client. +const ( + notifyQueueLimit = 4096 + notifyQueueByteLimit = 16 << 20 // 16 MiB +) + +// maxFrameBytes is the largest single frame this transport will read. Without +// it the byte budget above would not be authoritative: readMessage allocates the +// whole body before anything can classify it, so one notification with a huge +// Content-Length would be materialized in full and only then rejected by +// enqueueNotification. Capping at the same value keeps that impossible — a +// notification can never allocate past the budget it is measured against — and +// bounds what a malformed or hostile header can make this client allocate. A +// legitimate LSP frame is orders of magnitude smaller; one that is not is a +// protocol error, and failing the read closes the client so the manager restarts +// the session. +const maxFrameBytes = notifyQueueByteLimit + type rpcError struct { Code int `json:"code"` Message string `json:"message"` @@ -81,10 +139,14 @@ type incomingMessage struct { // server process exits); call Close to stop using the client. func NewClient(r io.Reader, w io.Writer) *Client { client := &Client{ - writer: w, - pending: make(map[int64]chan rpcResponse), - closed: make(chan struct{}), + writer: w, + pending: make(map[int64]chan rpcResponse), + closed: make(chan struct{}), + notifyReady: make(chan struct{}, 1), + notificationDrained: make(chan struct{}), + handledWait: make(chan struct{}), } + go client.notificationLoop() go client.readLoop(bufio.NewReader(r)) return client } @@ -155,6 +217,15 @@ func (c *Client) readLoop(reader *bufio.Reader) { c.failPending(err) return } + // Stamp the receipt boundary here — while the frame is nothing but bytes + // that have left the wire — rather than at enqueue. Everything below + // (json.Unmarshal of a peer-sized payload, then enqueue) runs on this + // goroutine while another can be capturing a baseline from ReceiptSeq, so + // stamping later would let a publish that was already read be numbered + // above a baseline taken after it arrived, and a superseded publish would + // then look fresh. Frames that turn out not to be notifications consume a + // number too: the sequence is a receipt clock, and gaps in it are fine. + seq := c.stampReceipt() var msg incomingMessage if err := json.Unmarshal(body, &msg); err != nil { continue // skip a malformed frame rather than tearing down the session @@ -166,12 +237,7 @@ func (c *Client) readLoop(reader *bufio.Reader) { // required or the server can block waiting on it (e.g. registerCapability). _ = c.write(outgoingReply{JSONRPC: "2.0", ID: msg.ID, Result: nil}) case msg.Method != "": - c.mu.Lock() - handler := c.handler - c.mu.Unlock() - if handler != nil { - handler(msg.Method, msg.Params) - } + c.enqueueNotification(notification{method: msg.Method, params: msg.Params, seq: seq}) case hasID: var id int64 if err := json.Unmarshal(msg.ID, &id); err == nil { @@ -181,6 +247,171 @@ func (c *Client) readLoop(reader *bufio.Reader) { } } +func (c *Client) notificationLoop() { + defer close(c.notificationDrained) + for { + <-c.notifyReady + for { + notification, ok, closed := c.dequeueNotification() + if !ok { + if closed { + return + } + break + } + c.mu.Lock() + handler := c.handler + c.mu.Unlock() + if handler != nil { + handler(notification.method, notification.params, notification.seq) + } + c.markHandled(notification.seq) + } + } +} + +// markHandled records that everything received through seq has now been handled +// and wakes anyone waiting for the worker to catch up. +func (c *Client) markHandled(seq int64) { + c.notifyMu.Lock() + if seq > c.handledSeq { + c.handledSeq = seq + } + wait := c.handledWait + c.handledWait = make(chan struct{}) + c.notifyMu.Unlock() + if wait != nil { + close(wait) + } +} + +// handledThrough returns the receipt seq the worker has finished handling, plus +// a channel closed the next time that advances. The channel is nil for a Client +// built without NewClient (some unit tests), which has no worker to advance it. +func (c *Client) handledThrough() (int64, <-chan struct{}) { + c.notifyMu.Lock() + defer c.notifyMu.Unlock() + return c.handledSeq, c.handledWait +} + +// acceptedNotificationSeq returns the receipt seq of the most recently queued +// notification (0 if none). It is the point the worker must reach to have +// handled everything the reader had accepted at the time of the call. +func (c *Client) acceptedNotificationSeq() int64 { + c.notifyMu.Lock() + defer c.notifyMu.Unlock() + return c.acceptedSeq +} + +// notificationsDone returns the worker completion signal. A nil result means +// this Client was constructed without NewClient (as some unit-test clients are) +// and has no notification worker to wait for. +func (c *Client) notificationsDone() <-chan struct{} { + return c.notificationDrained +} + +// enqueueNotification hands a server notification to the worker loop. item.seq +// must already carry the receipt stamp from stampReceipt: the receipt boundary +// belongs to the moment the frame left the wire, not to this call, which happens +// a full json.Unmarshal later. It never blocks and never silently discards a +// message the handler could still act on: the queue grows instead, up to its +// message and byte limits. +// +// The alternatives to growing are worse. Blocking the read loop when a buffer +// fills is the deadlock this dispatch exists to avoid — a handler that calls +// Client.Call waits for a response frame the blocked reader can no longer +// deliver. Dropping the oldest queued item instead loses protocol state +// permanently: a textDocument/publishDiagnostics for one URI is the server's +// only report for that URI, so discarding it makes session.waitForDiagnostics +// time out and Manager.Check return nothing even though the server published +// findings. +// +// Growth is bounded in practice by how much the server emits while a handler +// runs, and the queue is released as soon as it drains. But "in practice" is +// not a limit: a handler that never returns, or a server that sustains a +// higher rate than the single handler can drain, would otherwise grow this +// queue's full json.RawMessage payloads without bound until the heap gives +// out. The count limit alone is insufficient because notification payloads are +// peer-controlled and can be large; the byte limit prevents a few large +// diagnostics publishes from exhausting the heap before the count limit is +// reached. Either limit turns overload into an explicit, observable failure — +// the client is failed and closed — rather than unbounded protocol retention. +func (c *Client) enqueueNotification(item notification) { + c.notifyMu.Lock() + if c.notifyClosed { + // The worker loop has already stopped, so anything queued now would never + // be handled. Retaining it would grow the queue for as long as the + // transport stays readable after Close — Server.Shutdown closes the client + // before closing stdin, so a server emitting notifications while it handles + // shutdown/exit keeps the read loop feeding a queue nobody drains. + c.notifyMu.Unlock() + return + } + itemBytes := len(item.method) + len(item.params) + if len(c.notifyQueue) >= notifyQueueLimit || itemBytes > notifyQueueByteLimit-c.notifyBytes { + c.notifyMu.Unlock() + c.failPending(fmt.Errorf( + "lsp client: notification backlog exceeded %d messages or %d bytes", + notifyQueueLimit, + notifyQueueByteLimit, + )) + return + } + c.notifyQueue = append(c.notifyQueue, item) + c.notifyBytes += itemBytes + c.acceptedSeq = item.seq + c.notifyMu.Unlock() + + select { + case c.notifyReady <- struct{}{}: + default: + // A wake-up is already pending; the worker drains the whole queue per wake, + // so this item is covered by it. + } +} + +// stampReceipt claims the next receipt number for a frame that has just been +// read off the wire. The read loop calls it for every frame before parsing one, +// so a notification's seq reflects when it arrived rather than when this client +// got around to decoding and queueing it. +func (c *Client) stampReceipt() int64 { + c.notifyMu.Lock() + defer c.notifyMu.Unlock() + c.receiptSeq++ + return c.receiptSeq +} + +// ReceiptSeq returns how many frames have been read off the wire so far, +// including any notification still waiting to be dispatched to the handler. A +// caller that wants to know whether a notification newer than "now" has arrived +// should snapshot this before triggering whatever produces it, then require a +// subsequently-observed seq to be strictly greater: a notification whose frame +// had already been read at snapshot time has seq <= the snapshot, even if it is +// still queued and its handler doesn't run until afterward. +func (c *Client) ReceiptSeq() int64 { + c.notifyMu.Lock() + defer c.notifyMu.Unlock() + return c.receiptSeq +} + +func (c *Client) dequeueNotification() (notification, bool, bool) { + c.notifyMu.Lock() + defer c.notifyMu.Unlock() + if len(c.notifyQueue) == 0 { + return notification{}, false, c.notifyClosed + } + item := c.notifyQueue[0] + c.notifyBytes -= len(item.method) + len(item.params) + c.notifyQueue[0] = notification{} + c.notifyQueue = c.notifyQueue[1:] + if len(c.notifyQueue) == 0 { + // Release the backing array once drained; re-slicing alone would keep a + // burst's worth of capacity (and its already-consumed items) alive. + c.notifyQueue = nil + } + return item, true, false +} + func (c *Client) deliver(id int64, resp rpcResponse) { c.mu.Lock() ch, ok := c.pending[id] @@ -204,9 +435,27 @@ func (c *Client) failPending(err error) { ch <- rpcResponse{Err: &rpcError{Code: -1, Message: err.Error()}} } close(c.closed) + c.closeNotifications() }) } +// closeNotifications atomically stops accepting notifications, then wakes the +// worker so everything accepted before that boundary is delivered in FIFO order. +// It deliberately does not wait: a user handler may be blocked, and transport +// failure, overload, and concurrent Close must never deadlock teardown. The +// worker exits itself after the accepted queue (including any in-flight handler) +// has drained. An enqueue racing with shutdown is therefore either accepted and +// delivered before worker exit, or observes notifyClosed and is rejected. +func (c *Client) closeNotifications() { + c.notifyMu.Lock() + c.notifyClosed = true + c.notifyMu.Unlock() + select { + case c.notifyReady <- struct{}{}: + default: + } +} + func (c *Client) readError() error { c.mu.Lock() defer c.mu.Unlock() @@ -283,6 +532,9 @@ func readMessage(reader *bufio.Reader) ([]byte, error) { if contentLength < 0 { return nil, errors.New("message missing Content-Length header") } + if contentLength > maxFrameBytes { + return nil, fmt.Errorf("lsp frame of %d bytes exceeds the %d byte limit", contentLength, maxFrameBytes) + } body := make([]byte, contentLength) if _, err := io.ReadFull(reader, body); err != nil { return nil, err diff --git a/internal/lsp/client_test.go b/internal/lsp/client_test.go index 5ec6c39d0..7f4626b04 100644 --- a/internal/lsp/client_test.go +++ b/internal/lsp/client_test.go @@ -7,7 +7,9 @@ import ( "encoding/json" "fmt" "io" + "runtime" "strings" + "sync" "testing" "time" ) @@ -208,7 +210,7 @@ func TestClientNotificationHandler(t *testing.T) { _ = serverReader received := make(chan string, 1) - client.SetNotificationHandler(func(method string, _ json.RawMessage) { + client.SetNotificationHandler(func(method string, _ json.RawMessage, _ int64) { received <- method }) _ = writeMessage(serverWriter, map[string]any{ @@ -227,6 +229,438 @@ func TestClientNotificationHandler(t *testing.T) { } } +func TestClientNotificationHandlerCanCallClient(t *testing.T) { + clientReader, serverWriter := io.Pipe() + serverReader, clientWriter := io.Pipe() + client := NewClient(clientReader, clientWriter) + defer client.Close() + defer serverWriter.Close() + defer clientWriter.Close() + + serverDone := make(chan error, 1) + go func() { + body, err := readMessage(bufio.NewReader(serverReader)) + if err != nil { + serverDone <- err + return + } + var request incomingMessage + if err := json.Unmarshal(body, &request); err != nil { + serverDone <- err + return + } + serverDone <- writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", + "id": request.ID, + "result": map[string]bool{"applied": true}, + }) + }() + + handlerDone := make(chan error, 1) + client.SetNotificationHandler(func(_ string, _ json.RawMessage, _ int64) { + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err := client.Call(ctx, "workspace/applyEdit", nil) + handlerDone <- err + }) + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", + "method": "workspace/requestEdit", + }); err != nil { + t.Fatal(err) + } + + select { + case err := <-handlerDone: + if err != nil { + t.Fatalf("notification handler call failed: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("notification handler deadlocked waiting for its response") + } + if err := <-serverDone; err != nil { + t.Fatalf("server failed: %v", err) + } +} + +func TestClientNotificationHandlersPreserveOrder(t *testing.T) { + clientReader, serverWriter := io.Pipe() + client := NewClient(clientReader, io.Discard) + defer client.Close() + defer serverWriter.Close() + + received := make(chan string, 2) + client.SetNotificationHandler(func(method string, _ json.RawMessage, _ int64) { + received <- method + }) + for _, method := range []string{"first", "second"} { + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", + "method": method, + }); err != nil { + t.Fatal(err) + } + } + + for _, want := range []string{"first", "second"} { + select { + case got := <-received: + if got != want { + t.Fatalf("notification = %q, want %q", got, want) + } + case <-time.After(2 * time.Second): + t.Fatalf("timed out waiting for notification %q", want) + } + } +} + +func TestClientNotificationBurstDoesNotBlockResponse(t *testing.T) { + clientReader, serverWriter := io.Pipe() + serverReader, clientWriter := io.Pipe() + client := NewClient(clientReader, clientWriter) + defer client.Close() + defer serverWriter.Close() + defer clientWriter.Close() + + handlerStarted := make(chan struct{}) + handlerDone := make(chan error, 1) + var mu sync.Mutex + var queued []string + allQueued := make(chan struct{}) + client.SetNotificationHandler(func(method string, _ json.RawMessage, _ int64) { + if method != "blocking" { + mu.Lock() + queued = append(queued, method) + complete := len(queued) == notificationBurstSize + mu.Unlock() + if complete { + close(allQueued) + } + return + } + close(handlerStarted) + ctx, cancel := context.WithTimeout(context.Background(), time.Second) + defer cancel() + _, err := client.Call(ctx, "workspace/applyEdit", nil) + handlerDone <- err + }) + + serverDone := make(chan error, 1) + go func() { + reader := bufio.NewReader(serverReader) + body, err := readMessage(reader) + if err != nil { + serverDone <- err + return + } + var request incomingMessage + if err := json.Unmarshal(body, &request); err != nil { + serverDone <- err + return + } + for i := 0; i < notificationBurstSize; i++ { + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", + "method": fmt.Sprintf("queued-%03d", i), + }); err != nil { + serverDone <- err + return + } + } + serverDone <- writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", + "id": request.ID, + "result": nil, + }) + }() + + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", + "method": "blocking", + }); err != nil { + t.Fatal(err) + } + select { + case <-handlerStarted: + case <-time.After(time.Second): + t.Fatal("blocking notification handler did not start") + } + select { + case err := <-handlerDone: + if err != nil { + t.Fatalf("notification handler call failed under burst: %v", err) + } + case <-time.After(2 * time.Second): + t.Fatal("notification burst blocked the response read loop") + } + if err := <-serverDone; err != nil { + t.Fatalf("server failed: %v", err) + } + // Not blocking the reader is only half the requirement: every notification the + // server sent during the burst must still reach the handler, in order. A + // bounded queue that drops the oldest would silently lose the head of this + // sequence — and with it a publishDiagnostics the checker is waiting for. + select { + case <-allQueued: + case <-time.After(5 * time.Second): + mu.Lock() + received := len(queued) + mu.Unlock() + t.Fatalf("received %d of %d burst notifications; the queue lost messages", received, notificationBurstSize) + } + mu.Lock() + defer mu.Unlock() + for i, method := range queued { + if want := fmt.Sprintf("queued-%03d", i); method != want { + t.Fatalf("burst notification %d = %q, want %q (order must survive the burst)", i, method, want) + } + } +} + +// notificationBurstSize is far past any buffer the dispatch used to have, so a +// bounded queue would be visible as either a blocked read loop or a lost message. +const notificationBurstSize = 512 + +// TestClientNotificationQueueIsLossless is the regression test for the overflow +// policy: a queued notification must never be discarded. A dropped +// textDocument/publishDiagnostics is the server's only report for that URI, so +// losing it makes session.waitForDiagnostics time out and Manager.Check return +// nothing even though the server published findings. +func TestClientNotificationQueueIsLossless(t *testing.T) { + client := &Client{notifyReady: make(chan struct{}, 1)} + for i := 0; i < notificationBurstSize; i++ { + client.enqueueNotification(notification{method: fmt.Sprintf("notification-%03d", i)}) + } + + for i := 0; i < notificationBurstSize; i++ { + item, ok, _ := client.dequeueNotification() + if !ok { + t.Fatalf("notification %d was discarded by the queue", i) + } + want := fmt.Sprintf("notification-%03d", i) + if item.method != want { + t.Fatalf("notification = %q, want %q (queue must stay FIFO)", item.method, want) + } + } + if _, ok, _ := client.dequeueNotification(); ok { + t.Fatal("queue returned more notifications than were enqueued") + } + if client.notifyQueue != nil { + t.Fatalf("drained queue retained %d slots of capacity", cap(client.notifyQueue)) + } + if client.notifyBytes != 0 { + t.Fatalf("drained queue retained byte accounting: %d", client.notifyBytes) + } +} + +func TestClientDrainsAcceptedNotificationsAfterTransportEOF(t *testing.T) { + clientReader, serverWriter := io.Pipe() + client := NewClient(clientReader, io.Discard) + + blockStarted := make(chan struct{}) + releaseBlock := make(chan struct{}) + handled := make(chan string, 2) + client.SetNotificationHandler(func(method string, _ json.RawMessage, _ int64) { + if method == "block" { + close(blockStarted) + <-releaseBlock + } + handled <- method + }) + + if err := writeMessage(serverWriter, map[string]any{"jsonrpc": "2.0", "method": "block"}); err != nil { + t.Fatal(err) + } + <-blockStarted + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", "method": "textDocument/publishDiagnostics", + }); err != nil { + t.Fatal(err) + } + // Wait until the reader has accepted the publish, then fail it while the + // sole worker is still in the preceding handler and cannot dequeue it. + deadline := time.Now().Add(time.Second) + for client.ReceiptSeq() != 2 { + if time.Now().After(deadline) { + t.Fatal("publish was not accepted by the read loop") + } + runtime.Gosched() + } + if err := serverWriter.Close(); err != nil { + t.Fatal(err) + } + // Fresh budget: the wait above may have consumed most of the previous one, + // which would make this loop fail on nothing worse than slow scheduling. + deadline = time.Now().Add(time.Second) + for !client.IsClosed() { + if time.Now().After(deadline) { + t.Fatal("transport EOF did not close client") + } + runtime.Gosched() + } + close(releaseBlock) + + for _, want := range []string{"block", "textDocument/publishDiagnostics"} { + select { + case got := <-handled: + if got != want { + t.Fatalf("handled %q, want %q", got, want) + } + case <-time.After(time.Second): + t.Fatalf("accepted notification %q was not delivered after EOF", want) + } + } + select { + case <-client.notificationDrained: + case <-time.After(time.Second): + t.Fatal("notification worker did not report the accepted queue drained") + } +} + +// TestReceiptStampedAtReadNotEnqueue is the regression test for the review gap +// on #759 (pullrequestreview-4834747507): TestPublishBaselineRejectsFrameReadBeforeBaseline +// calls stampReceipt itself and never drives readLoop, so it pins the seq +// comparison but not WHERE the stamp is taken — a refactor that moved stamping +// back to the enqueue call site (reverting the read-time fix) left that test +// green. +// +// This test distinguishes the two call sites directly: a response frame is +// read off the wire but never reaches enqueueNotification (see the `hasID` +// branch in readLoop). Under read-time stamping every frame consumes a receipt +// number — "gaps in it are fine", per stampReceipt's doc comment — so +// ReceiptSeq must advance even though nothing was queued. Under enqueue-time +// stamping it would not, since stampReceipt would never run for this frame. +func TestReceiptStampedAtReadNotEnqueue(t *testing.T) { + clientReader, serverWriter := io.Pipe() + client := NewClient(clientReader, io.Discard) + defer serverWriter.Close() + + // A plain response frame: has an id, no method, so it is delivered via + // c.deliver and never passed to enqueueNotification. + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", "id": 1, "result": nil, + }); err != nil { + t.Fatal(err) + } + + deadline := time.Now().Add(time.Second) + for client.ReceiptSeq() == 0 { + if time.Now().After(deadline) { + t.Fatal("ReceiptSeq did not advance for a frame that was read but never enqueued as a notification") + } + runtime.Gosched() + } +} + +// TestReadMessageRejectsOversizedFrame covers jatmn's P3 finding on #759: the +// notification byte budget was enforced only at enqueue, but readMessage +// allocates the whole body before anything can tell a notification from a +// response. A single frame declaring a huge Content-Length was therefore +// materialized in full and only then measured against a budget it had already +// blown past. The header is now rejected before the allocation. +func TestReadMessageRejectsOversizedFrame(t *testing.T) { + header := fmt.Sprintf("Content-Length: %d\r\n\r\n", maxFrameBytes+1) + _, err := readMessage(bufio.NewReader(strings.NewReader(header))) + if err == nil { + t.Fatal("an oversized frame must be rejected, not allocated and read") + } + if !strings.Contains(err.Error(), "exceeds") { + t.Fatalf("error = %v, want the frame-size limit to reject it before reading the body", err) + } + // A frame at the limit is still legal: the cap rejects what is over it, not + // what merely approaches it. + body := strings.Repeat("x", 16) + atLimit := fmt.Sprintf("Content-Length: %d\r\n\r\n%s", len(body), body) + if _, err := readMessage(bufio.NewReader(strings.NewReader(atLimit))); err != nil { + t.Fatalf("readMessage of an ordinary frame: %v", err) + } +} + +// TestClientFailsOnNotificationBacklogOverload is the regression test for +// jatmn's #759 P2 finding: the lossless queue above had no failure policy — a +// language server sustaining a higher notification rate than the single +// handler can drain (no permanently-stuck handler required, just a sustained +// producer faster than the consumer) would retain every full json.RawMessage +// payload on Zero's heap without bound. Hitting notifyQueueLimit must fail +// (and close) the client observably instead of continuing to grow. +func TestClientFailsOnNotificationBacklogOverload(t *testing.T) { + client := &Client{ + notifyReady: make(chan struct{}, 1), + pending: make(map[int64]chan rpcResponse), + closed: make(chan struct{}), + } + for i := 0; i < notifyQueueLimit; i++ { + client.enqueueNotification(notification{method: "spam"}) + } + if client.IsClosed() { + t.Fatal("client closed before the backlog limit was reached") + } + client.notifyMu.Lock() + queued := len(client.notifyQueue) + client.notifyMu.Unlock() + if queued != notifyQueueLimit { + t.Fatalf("queued = %d, want %d before the limit is exceeded", queued, notifyQueueLimit) + } + + // One push past the limit must fail the client rather than growing the + // queue further. + client.enqueueNotification(notification{method: "spam"}) + if !client.IsClosed() { + t.Fatal("client must be closed once the notification backlog exceeds notifyQueueLimit") + } + client.notifyMu.Lock() + queuedAfter := len(client.notifyQueue) + bytesAfter := client.notifyBytes + client.notifyMu.Unlock() + if queuedAfter != notifyQueueLimit { + t.Fatalf("closed client retained %d accepted notifications, want %d pending worker drain", queuedAfter, notifyQueueLimit) + } + if bytesAfter == 0 { + t.Fatal("accepted notifications were discarded instead of left for worker drain") + } +} + +func TestClientFailsOnNotificationBacklogByteOverload(t *testing.T) { + client := &Client{ + notifyReady: make(chan struct{}, 1), + pending: make(map[int64]chan rpcResponse), + closed: make(chan struct{}), + } + const messageCount = 16 + const method = "spam" + for i := 0; i < messageCount; i++ { + client.enqueueNotification(notification{ + method: method, + params: make(json.RawMessage, notifyQueueByteLimit/messageCount-len(method)), + }) + } + if client.IsClosed() { + t.Fatal("client closed before the notification byte limit was exceeded") + } + client.notifyMu.Lock() + queued := len(client.notifyQueue) + queuedBytes := client.notifyBytes + client.notifyMu.Unlock() + if queued != messageCount { + t.Fatalf("queued = %d, want %d at the byte limit", queued, messageCount) + } + if queuedBytes != notifyQueueByteLimit { + t.Fatalf("queued bytes = %d, want limit %d", queuedBytes, notifyQueueByteLimit) + } + + client.enqueueNotification(notification{method: "x"}) + if !client.IsClosed() { + t.Fatal("client must close before retaining notification bytes past the limit") + } + client.notifyMu.Lock() + queuedAfter := len(client.notifyQueue) + bytesAfter := client.notifyBytes + client.notifyMu.Unlock() + if queuedAfter != messageCount || bytesAfter != notifyQueueByteLimit { + t.Fatalf("accepted backlog changed during close: %d notifications and %d accounted bytes", queuedAfter, bytesAfter) + } +} + func TestClientRejectsCallsAfterClose(t *testing.T) { clientReader, serverWriter := io.Pipe() serverReader, clientWriter := io.Pipe() @@ -243,3 +677,79 @@ func TestClientRejectsCallsAfterClose(t *testing.T) { t.Fatal("Notify after Close must return an error") } } + +// TestClientDropsNotificationsAfterClose covers the shutdown path: Close stops +// the worker loop, but the read loop keeps reading until the transport ends — +// Server.Shutdown closes the client BEFORE closing stdin, so a server emitting +// notifications while it handles shutdown/exit would otherwise pile them into a +// queue nobody drains. Nothing may be handled or retained after Close. +func TestClientDropsNotificationsAfterClose(t *testing.T) { + clientReader, serverWriter := io.Pipe() + serverReader, clientWriter := io.Pipe() + client := NewClient(clientReader, clientWriter) + defer serverWriter.Close() + defer clientWriter.Close() + go func() { + // Drain anything the client writes so no write can block the test. + _, _ = io.Copy(io.Discard, serverReader) + }() + + handled := make(chan string, 4) + client.SetNotificationHandler(func(method string, _ json.RawMessage, _ int64) { + handled <- method + }) + + // A notification before Close still reaches the handler. + if err := writeMessage(serverWriter, map[string]any{"jsonrpc": "2.0", "method": "before-close"}); err != nil { + t.Fatal(err) + } + select { + case method := <-handled: + if method != "before-close" { + t.Fatalf("handled %q, want before-close", method) + } + case <-time.After(2 * time.Second): + t.Fatal("notification before Close was not handled") + } + + if err := client.Close(); err != nil { + t.Fatalf("Close: %v", err) + } + // The transport stays writable, exactly as it does between Client.Close and + // stdin.Close in Server.Shutdown. + for i := 0; i < 64; i++ { + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", + "method": fmt.Sprintf("after-close-%03d", i), + }); err != nil { + t.Fatalf("write notification after Close: %v", err) + } + } + + select { + case method := <-handled: + t.Fatalf("handler ran for %q after Close", method) + case <-time.After(200 * time.Millisecond): + } + + // Give the read loop time to consume every frame it was handed, then require + // the queue to hold nothing: dropping on enqueue is what keeps a closed client + // from growing for as long as its reader stays open. + deadline := time.Now().Add(2 * time.Second) + for { + client.notifyMu.Lock() + queued := len(client.notifyQueue) + closed := client.notifyClosed + client.notifyMu.Unlock() + if !closed { + t.Fatal("Close did not mark the notification queue closed") + } + if queued != 0 { + t.Fatalf("closed client retained %d queued notifications", queued) + } + if time.Now().After(deadline) { + return + } + time.Sleep(10 * time.Millisecond) + } +} diff --git a/internal/lsp/documents.go b/internal/lsp/documents.go index fa7d1ec63..589517f47 100644 --- a/internal/lsp/documents.go +++ b/internal/lsp/documents.go @@ -20,34 +20,39 @@ type lspServer interface { type session struct { server lspServer client *Client + // catchUpGrace bounds catchUpNotifications on a live client. Zero means "do + // not wait", which is what the hand-built sessions in tests get by default. + catchUpGrace time.Duration - mu sync.Mutex - open map[string]bool // uri -> didOpen sent - versions map[string]int // uri -> current (committed) version - diagnostics map[string][]Diagnostic // uri -> latest published diagnostics - lastPublish map[string]time.Time // uri -> time of last publish - publishCount map[string]int // uri -> monotonic publish count - waiters map[string][]chan struct{} // uri -> goroutines awaiting the next publish - fileLocks map[string]*sync.Mutex // uri -> per-document sync serializer + mu sync.Mutex + open map[string]bool // uri -> didOpen sent + versions map[string]int // uri -> current (committed) version + diagnostics map[string][]Diagnostic // uri -> latest published diagnostics + lastPublish map[string]time.Time // uri -> time of last publish + publishSeq map[string]int64 // uri -> receipt seq of latest publish (see Client.ReceiptSeq) + waiters map[string][]chan struct{} // uri -> goroutines awaiting the next publish + fileLocks map[string]*sync.Mutex // uri -> per-document sync serializer } func newSession(server lspServer) *session { s := &session{ server: server, client: server.Client(), - open: map[string]bool{}, - versions: map[string]int{}, - diagnostics: map[string][]Diagnostic{}, - lastPublish: map[string]time.Time{}, - publishCount: map[string]int{}, - waiters: map[string][]chan struct{}{}, - fileLocks: map[string]*sync.Mutex{}, + catchUpGrace: notificationCatchUpGrace, + + open: map[string]bool{}, + versions: map[string]int{}, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + fileLocks: map[string]*sync.Mutex{}, } s.client.SetNotificationHandler(s.handleNotification) return s } -func (s *session) handleNotification(method string, params json.RawMessage) { +func (s *session) handleNotification(method string, params json.RawMessage, seq int64) { if method != "textDocument/publishDiagnostics" { return } @@ -66,7 +71,7 @@ func (s *session) handleNotification(method string, params json.RawMessage) { } s.diagnostics[payload.URI] = payload.Diagnostics s.lastPublish[payload.URI] = time.Now() - s.publishCount[payload.URI]++ + s.publishSeq[payload.URI] = seq waiters := s.waiters[payload.URI] delete(s.waiters, payload.URI) s.mu.Unlock() @@ -133,36 +138,65 @@ func (s *session) diagnosticsFor(uri string) []Diagnostic { return append([]Diagnostic(nil), s.diagnostics[uri]...) } -// publishBaseline snapshots the current publish count for a URI, captured before -// a sync so waitForDiagnostics can wait specifically for the publish that sync -// triggers (not be satisfied by a stale earlier publish). -func (s *session) publishBaseline(uri string) int { - s.mu.Lock() - defer s.mu.Unlock() - return s.publishCount[uri] +// publishBaseline snapshots the client's current notification receipt +// sequence, captured before a sync so waitForDiagnostics can require a publish +// that was RECEIVED after this point. This must baseline against receipt, not +// against session.publishSeq (which only advances once handleNotification +// actually runs): dispatch happens off a queue, so a publish for a since- +// superseded version can already be sitting in that queue, still unprocessed, +// at the moment a later Check captures its baseline. If baseline were instead +// "how many publishes has this session recorded so far", that stale publish +// would land afterward, still satisfy "more than baseline", and hand the new +// check diagnostics for the wrong text — silently, since many servers omit the +// version field the staleness check in handleNotification relies on. +// Baselining against receipt sequence closes that: a notification whose frame +// had already come off the wire before this call has seq <= baseline no matter +// when it is decoded, queued, or handled. +func (s *session) publishBaseline() int64 { + return s.client.ReceiptSeq() } +// notificationCatchUpGrace is the live-client catch-up budget a real session +// runs with. It exists for the case where the caller's own deadline has already +// passed: dropping a publish that is sitting accepted-but-unhandled in the queue +// would report "no diagnostics" for text the server had in fact already answered +// for, so a brief overrun is preferable to discarding it. It stays short because +// the caller is by then out of budget, and it costs nothing in the ordinary case +// — the session's own handler only parses and records, so the worker is +// virtually always caught up already and no wait happens at all. A closed client +// is deliberately not bounded by this; see catchUpNotifications. +const notificationCatchUpGrace = 50 * time.Millisecond + // waitForDiagnostics blocks until a publish newer than baseline arrives for the // URI and the server then goes quiet for the debounce window, or until ctx is -// done. It returns true when a fresh publish was observed and false when ctx -// expired before any did — so a caller can avoid returning a stale prior result -// for newer text. Servers don't signal "analysis complete", so the debounce -// approximates it: once a fresh publish lands, wait debounce for a follow-up, -// resetting on each new publish. -func (s *session) waitForDiagnostics(ctx context.Context, uri string, debounce time.Duration, baseline int) bool { +// done or the client closes. It returns true when a fresh publish was observed +// and false when the wait ended before any did — so a caller can avoid returning +// a stale prior result for newer text. Servers don't signal "analysis complete", +// so the debounce approximates it: once a fresh publish lands, wait debounce for +// a follow-up, resetting on each new publish. +func (s *session) waitForDiagnostics(ctx context.Context, uri string, debounce time.Duration, baseline int64) bool { for { s.mu.Lock() ch := make(chan struct{}) s.waiters[uri] = append(s.waiters[uri], ch) - count := s.publishCount[uri] + seq := s.publishSeq[uri] last := s.lastPublish[uri] s.mu.Unlock() - if count <= baseline { + if seq <= baseline { select { case <-ctx.Done(): s.cancelWaiter(uri, ch) - return false // no fresh publish for this sync arrived in time + // Cancellation, like closure below, ends the wait on a signal that + // says nothing about what the server sent, and either can win a race + // against a publish this session has already received. Deciding + // "nothing arrived" straight out of the select would discard it, so + // both settle identically: let the worker finish what it accepted, + // then re-read what got recorded. + return s.freshAfterCatchUp(uri, baseline) + case <-s.client.closed: + s.cancelWaiter(uri, ch) + return s.freshAfterCatchUp(uri, baseline) case <-ch: continue // a fresh publish arrived; loop into the debounce check } @@ -171,6 +205,7 @@ func (s *session) waitForDiagnostics(ctx context.Context, uri string, debounce t remaining := debounce - time.Since(last) if remaining <= 0 { s.cancelWaiter(uri, ch) + s.catchUpNotifications() return true } timer := time.NewTimer(remaining) @@ -178,17 +213,83 @@ func (s *session) waitForDiagnostics(ctx context.Context, uri string, debounce t case <-ctx.Done(): timer.Stop() s.cancelWaiter(uri, ch) - return true // a fresh publish did arrive; ctx merely cut the debounce short + // A fresh publish did arrive; ctx merely cut the debounce short. Catch + // the worker up regardless: the caller reads diagnosticsFor the moment + // this returns, and a newer publish for this URI may be sitting accepted + // but unhandled — it should get the newest received, not the older one + // that happened to wake this wait. + s.catchUpNotifications() + return true + case <-s.client.closed: + timer.Stop() + s.cancelWaiter(uri, ch) + // Same, for whatever the server managed to publish before it died. + s.catchUpNotifications() + return true // preserve the fresh publish already received before failure case <-ch: timer.Stop() continue // a newer publish arrived; re-arm the debounce case <-timer.C: s.cancelWaiter(uri, ch) + s.catchUpNotifications() return true // quiet for the full window } } } +// freshAfterCatchUp reports whether a publish received after baseline has been +// recorded for the URI, having first let the notification worker finish frames +// it already accepted. Receipt is what the answer turns on: a publish that was +// read off the wire but is still queued would otherwise read as "never arrived". +func (s *session) freshAfterCatchUp(uri string, baseline int64) bool { + s.catchUpNotifications() + s.mu.Lock() + defer s.mu.Unlock() + return s.publishSeq[uri] > baseline +} + +// catchUpNotifications waits for the client's notification worker to handle +// every notification it had already accepted, so a decision made right after +// cannot miss a publish that was on the wire before it. It returns immediately +// in the ordinary case, where the worker is already caught up. +func (s *session) catchUpNotifications() { + drained := s.client.notificationsDone() + if drained == nil { + return // hand-built client (some unit tests): no worker to catch up on + } + if s.client.IsClosed() { + // Wait the worker out instead of time-boxing it. Admission is already + // closed so the backlog cannot grow, and a handler's re-entrant Call now + // fails immediately rather than blocking on a response that will never + // come — the drain is bounded, and it is the last chance these frames get. + <-drained + return + } + timer := time.NewTimer(s.catchUpGrace) + defer timer.Stop() + for { + // Re-read the target every pass rather than snapshotting it once. The + // reader keeps running while this waits, so a snapshot goes stale: the + // worker finishing the item that was last when we started satisfies + // `handled >= target` while a newer publish for the same URI sits queued, + // and the caller then answers from diagnosticsFor with the older one — + // the exact failure the closed-client drain exists to prevent, on the + // live path. Growth is not unbounded: the grace below caps the wait no + // matter how fast the server publishes. + handled, advanced := s.client.handledThrough() + if handled >= s.client.acceptedNotificationSeq() { + return + } + select { + case <-advanced: + case <-drained: + return + case <-timer.C: + return + } + } +} + // cancelWaiter removes a still-open waiter (one a publish has not already closed // and cleared) so it can't leak or be closed twice. func (s *session) cancelWaiter(uri string, target chan struct{}) { diff --git a/internal/lsp/manager.go b/internal/lsp/manager.go index 66d811bef..ef2bce6f3 100644 --- a/internal/lsp/manager.go +++ b/internal/lsp/manager.go @@ -73,7 +73,7 @@ func (m *Manager) Check(ctx context.Context, path, text string) ([]Diagnostic, e } return nil, err } - baseline := sess.publishBaseline(uri) + baseline := sess.publishBaseline() if err := sess.sync(ctx, abs, languageID, text); err != nil { return nil, err } diff --git a/internal/lsp/manager_test.go b/internal/lsp/manager_test.go index e646de138..a096b71d5 100644 --- a/internal/lsp/manager_test.go +++ b/internal/lsp/manager_test.go @@ -7,6 +7,7 @@ import ( "errors" "io" "os/exec" + "runtime" "strings" "sync" "testing" @@ -228,18 +229,748 @@ func TestSessionDropsStaleVersionPublish(t *testing.T) { sess.mu.Unlock() stale, _ := json.Marshal(PublishDiagnosticsParams{URI: uri, Version: 2, Diagnostics: []Diagnostic{{Message: "stale"}}}) - sess.handleNotification("textDocument/publishDiagnostics", stale) + sess.handleNotification("textDocument/publishDiagnostics", stale, 1) if len(sess.diagnosticsFor(uri)) != 0 { t.Fatal("a publish for an older version must be ignored") } fresh, _ := json.Marshal(PublishDiagnosticsParams{URI: uri, Version: 3, Diagnostics: []Diagnostic{{Message: "fresh"}}}) - sess.handleNotification("textDocument/publishDiagnostics", fresh) + sess.handleNotification("textDocument/publishDiagnostics", fresh, 2) if d := sess.diagnosticsFor(uri); len(d) != 1 || d[0].Message != "fresh" { t.Fatalf("a current-version publish must apply, got %#v", d) } } +// TestPublishBaselineRejectsAlreadyQueuedPublish is the regression test for +// jatmn's #759 P1 finding: publishBaseline used to snapshot how many publishes +// session.handleNotification had already RUN for a URI. Dispatch happens off a +// queue, though, so a publish for a version Check is about to supersede can +// already be sitting RECEIVED but undispatched at the moment a later Check +// captures its baseline — then run only afterward, incrementing the old +// count-based baseline and wrongly looking "new". Since many servers omit the +// version field, handleNotification's own staleness check (which only fires +// for a positive version) can't catch this either, so the stale result would +// reach the caller for the new text, and the debounce could finish before the +// real response even arrives. +// +// This drives the exact interleaving through the real production path: the +// stale publish is enqueued first (fixing its receipt seq), THEN a baseline is +// captured, THEN the stale item is dispatched — reproducing "received before +// baseline, handled after" without depending on goroutine scheduling luck. A +// bare Client (no background loops, like TestClientNotificationQueueIsLossless +// uses) makes the enqueue/dequeue ordering fully explicit. +func TestPublishBaselineRejectsAlreadyQueuedPublish(t *testing.T) { + client := &Client{notifyReady: make(chan struct{}, 1)} + sess := &session{ + client: client, + versions: map[string]int{}, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + uri := PathToURI("/repo/main.go") + + // The stale (version-less — the common case) publish is RECEIVED first: the + // read loop stamps a frame's receipt before it decodes or queues it, so the + // stamp here comes first too. + stale, _ := json.Marshal(PublishDiagnosticsParams{URI: uri, Diagnostics: []Diagnostic{{Message: "stale"}}}) + client.enqueueNotification(notification{ + method: "textDocument/publishDiagnostics", params: stale, seq: client.stampReceipt(), + }) + + // A later Check captures its baseline only now — after the stale publish's + // receipt, exactly as publishBaseline does between two Checks whose + // notification queue has backed up. + baseline := sess.publishBaseline() + + // Dispatch the stale item exactly as notificationLoop would: dequeue and + // call the handler with the receipt seq it was actually stamped with. + item, ok, _ := client.dequeueNotification() + if !ok { + t.Fatal("stale notification was not queued") + } + sess.handleNotification(item.method, item.params, item.seq) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + if sess.waitForDiagnostics(ctx, uri, 10*time.Millisecond, baseline) { + t.Fatal("a publish received before baseline must not satisfy waitForDiagnostics merely because it was handled afterward") + } + + // The real response — received (and handled) after baseline — must satisfy it. + fresh, _ := json.Marshal(PublishDiagnosticsParams{URI: uri, Diagnostics: []Diagnostic{{Message: "fresh"}}}) + client.enqueueNotification(notification{ + method: "textDocument/publishDiagnostics", params: fresh, seq: client.stampReceipt(), + }) + item2, ok, _ := client.dequeueNotification() + if !ok { + t.Fatal("fresh notification was not queued") + } + sess.handleNotification(item2.method, item2.params, item2.seq) + + ctx2, cancel2 := context.WithTimeout(context.Background(), time.Second) + defer cancel2() + if !sess.waitForDiagnostics(ctx2, uri, 10*time.Millisecond, baseline) { + t.Fatal("a publish received after baseline must satisfy waitForDiagnostics") + } + if d := sess.diagnosticsFor(uri); len(d) != 1 || d[0].Message != "fresh" { + t.Fatalf("diagnostics = %#v, want the fresh publish", d) + } +} + +// TestPublishBaselineRejectsFrameReadBeforeBaseline is the regression test for +// jatmn's follow-up P1 finding on #759: stamping receipt at ENQUEUE left a +// window async dispatch had opened. The read loop consumes a frame, then decodes +// it (a full json.Unmarshal of a peer-sized payload) and only then enqueues it, +// all while another goroutine can capture a baseline. A publishDiagnostics for +// text the new Check is about to supersede could therefore be numbered above a +// baseline taken after it had already come off the wire, and — since most +// servers omit the version field that handleNotification's staleness check needs +// — reach the caller as if it answered the new text. +// +// The interleaving is driven at the seam itself: stamp (the frame leaves the +// wire), baseline (a concurrent Check), then enqueue and dispatch (the read loop +// finishes decoding). Stamping at enqueue put this publish above the baseline; +// stamping at read puts it at or below. +func TestPublishBaselineRejectsFrameReadBeforeBaseline(t *testing.T) { + client := &Client{notifyReady: make(chan struct{}, 1)} + sess := &session{ + client: client, + versions: map[string]int{}, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + uri := PathToURI("/repo/main.go") + + // The read loop pulls the stale frame off the wire and stamps it... + seq := client.stampReceipt() + + // ...a concurrent Check baselines while that frame is still being decoded... + baseline := sess.publishBaseline() + + // ...and only now does the frame reach the queue and the handler. + stale, _ := json.Marshal(PublishDiagnosticsParams{URI: uri, Diagnostics: []Diagnostic{{Message: "stale"}}}) + client.enqueueNotification(notification{ + method: "textDocument/publishDiagnostics", params: stale, seq: seq, + }) + item, ok, _ := client.dequeueNotification() + if !ok { + t.Fatal("stale notification was not queued") + } + sess.handleNotification(item.method, item.params, item.seq) + + ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond) + defer cancel() + if sess.waitForDiagnostics(ctx, uri, 10*time.Millisecond, baseline) { + t.Fatal("a publish whose frame was read before baseline must not satisfy waitForDiagnostics merely because it was decoded and queued afterward") + } + if d := sess.diagnosticsFor(uri); len(d) != 1 || d[0].Message != "stale" { + t.Fatalf("diagnostics = %#v, want the stale publish still recorded (rejected as an answer, not dropped)", d) + } +} + +func TestWaitForDiagnosticsReturnsWhenClientCloses(t *testing.T) { + client := &Client{ + pending: make(map[int64]chan rpcResponse), + closed: make(chan struct{}), + } + uri := PathToURI("/repo/main.go") + sess := &session{ + client: client, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + waitDone := make(chan bool, 1) + go func() { + waitDone <- sess.waitForDiagnostics(context.Background(), uri, time.Second, 0) + }() + + deadline := time.Now().Add(time.Second) + for { + sess.mu.Lock() + waiting := len(sess.waiters[uri]) == 1 + sess.mu.Unlock() + if waiting { + break + } + if time.Now().After(deadline) { + t.Fatal("diagnostic wait was not registered") + } + time.Sleep(time.Millisecond) + } + + client.failPending(errors.New("notification backlog exceeded")) + select { + case fresh := <-waitDone: + if fresh { + t.Fatal("client failure without a fresh publish must not report fresh diagnostics") + } + case <-time.After(time.Second): + t.Fatal("diagnostic wait did not wake when the client failed") + } + sess.mu.Lock() + waiters := len(sess.waiters[uri]) + sess.mu.Unlock() + if waiters != 0 { + t.Fatalf("client failure left %d diagnostic waiters registered", waiters) + } +} + +func TestWaitForDiagnosticsPreservesPublishHandledBeforeClientCloses(t *testing.T) { + // Keep the test goroutine running while it closes both signals so the waiter + // observes the publish and client closure as simultaneously ready. Without + // re-checking publishSeq, select could randomly choose closure and discard the + // fresh diagnostics that handleNotification had already recorded. + previousProcs := runtime.GOMAXPROCS(1) + defer runtime.GOMAXPROCS(previousProcs) + + params, err := json.Marshal(PublishDiagnosticsParams{ + URI: PathToURI("/repo/main.go"), + Diagnostics: []Diagnostic{{Message: "fresh"}}, + }) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 100; i++ { + client := &Client{ + pending: make(map[int64]chan rpcResponse), + closed: make(chan struct{}), + } + uri := PathToURI("/repo/main.go") + sess := &session{ + client: client, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + waitDone := make(chan bool, 1) + go func() { + waitDone <- sess.waitForDiagnostics(context.Background(), uri, time.Second, 0) + }() + + deadline := time.Now().Add(time.Second) + for { + sess.mu.Lock() + waiting := len(sess.waiters[uri]) == 1 + sess.mu.Unlock() + if waiting { + break + } + if time.Now().After(deadline) { + t.Fatal("diagnostic wait was not registered") + } + time.Sleep(time.Millisecond) + } + + sess.handleNotification("textDocument/publishDiagnostics", params, 1) + client.failPending(errors.New("server exited after publishing")) + select { + case fresh := <-waitDone: + if !fresh { + t.Fatalf("iteration %d discarded diagnostics handled before client closure", i) + } + case <-time.After(time.Second): + t.Fatal("diagnostic wait did not wake when the client failed") + } + } +} + +func TestWaitForDiagnosticsDrainsAcceptedPublishAfterTransportEOF(t *testing.T) { + clientReader, serverWriter := io.Pipe() + client := NewClient(clientReader, io.Discard) + uri := PathToURI("/repo/main.go") + sess := &session{ + client: client, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + blocked := make(chan struct{}) + release := make(chan struct{}) + client.SetNotificationHandler(func(method string, params json.RawMessage, seq int64) { + if method == "test/block" { + close(blocked) + <-release + return + } + sess.handleNotification(method, params, seq) + }) + + if err := writeMessage(serverWriter, map[string]any{"jsonrpc": "2.0", "method": "test/block"}); err != nil { + t.Fatal(err) + } + <-blocked + params := PublishDiagnosticsParams{URI: uri, Diagnostics: []Diagnostic{{Message: "fresh"}}} + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", "method": "textDocument/publishDiagnostics", "params": params, + }); err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(time.Second) + for client.ReceiptSeq() != 2 { + if time.Now().After(deadline) { + t.Fatal("publish was not accepted before EOF") + } + runtime.Gosched() + } + + waitDone := make(chan bool, 1) + go func() { waitDone <- sess.waitForDiagnostics(context.Background(), uri, 0, 0) }() + if err := serverWriter.Close(); err != nil { + t.Fatal(err) + } + // Fresh budget: the wait above may have consumed most of the previous one, + // which would make this loop fail on nothing worse than slow scheduling. + deadline = time.Now().Add(time.Second) + for !client.IsClosed() { + if time.Now().After(deadline) { + t.Fatal("transport EOF did not close client") + } + runtime.Gosched() + } + select { + case fresh := <-waitDone: + t.Fatalf("wait returned %v before accepted notifications drained", fresh) + case <-time.After(25 * time.Millisecond): + } + close(release) + select { + case fresh := <-waitDone: + if !fresh { + t.Fatal("wait missed publish accepted before transport EOF") + } + case <-time.After(time.Second): + t.Fatal("wait did not finish after notification drain") + } + if got := sess.diagnosticsFor(uri); len(got) != 1 || got[0].Message != "fresh" { + t.Fatalf("diagnostics = %#v, want fresh publish", got) + } +} + +// TestWaitForDiagnosticsCatchesUpOnContextCancellation covers jatmn's P2 finding +// (1): the cancellation arm used to re-check only what the handler had already +// RUN. A publishDiagnostics accepted off the wire but still queued behind a busy +// worker therefore read as "never arrived", and Manager.Check returned (nil, +// nil) for text the server had in fact answered. Cancellation now settles like +// closure does — let the worker finish what it accepted, then re-read. +// TestCatchUpNotificationsFollowsPublishesAcceptedWhileWaiting covers jatmn's +// #759 finding: catch-up snapshotted the accepted sequence once, so a publish +// the reader accepted DURING the wait was not waited for. The worker finishing +// the stale target satisfied the check, and the caller then read diagnosticsFor +// while a newer publish for the same URI was still queued — the same +// "newest received, not the one that woke the wait" failure the closed-client +// drain regression was written for, on the live path. +func TestCatchUpNotificationsFollowsPublishesAcceptedWhileWaiting(t *testing.T) { + clientReader, serverWriter := io.Pipe() + defer serverWriter.Close() + client := NewClient(clientReader, io.Discard) + uri := PathToURI("/repo/main.go") + sess := &session{ + client: client, + catchUpGrace: 5 * time.Second, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + + publish := func(message string) { + t.Helper() + params := PublishDiagnosticsParams{URI: uri, Diagnostics: []Diagnostic{{Message: message}}} + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", "method": "textDocument/publishDiagnostics", "params": params, + }); err != nil { + t.Error(err) + } + } + + blocked := make(chan struct{}) + release := make(chan struct{}) + secondEntered := make(chan struct{}) + releaseSecond := make(chan struct{}) + client.SetNotificationHandler(func(method string, params json.RawMessage, seq int64) { + switch { + case method == "test/block": + close(blocked) + <-release + return + case strings.Contains(string(params), `"first"`): + // Handling "first" is what makes "second" arrive mid-catch-up: it goes + // on the wire and is confirmed accepted before this handler returns, so + // the worker's handled count only reaches "first" once a newer item is + // already queued behind it. + publish("second") + deadline := time.Now().Add(time.Second) + for client.acceptedNotificationSeq() < 3 { + if time.Now().After(deadline) { + t.Error("follow-up publish was not accepted while the worker was busy") + break + } + runtime.Gosched() + } + case strings.Contains(string(params), `"second"`): + // Hold the worker inside this handler so the assertion below observes + // the state the bug produced: caught up by the stale target, with the + // newer publish still unrecorded. + close(secondEntered) + <-releaseSecond + } + sess.handleNotification(method, params, seq) + }) + + if err := writeMessage(serverWriter, map[string]any{"jsonrpc": "2.0", "method": "test/block"}); err != nil { + t.Fatal(err) + } + <-blocked + publish("first") + deadline := time.Now().Add(time.Second) + for client.ReceiptSeq() != 2 { + if time.Now().After(deadline) { + t.Fatal("first publish was not accepted") + } + runtime.Gosched() + } + + done := make(chan struct{}) + go func() { + sess.catchUpNotifications() + close(done) + }() + // Catch-up must be parked in its wait before the worker is released, so the + // sequence it started from is the one that goes stale. + select { + case <-done: + t.Fatal("catch-up returned while the worker was still behind") + case <-time.After(25 * time.Millisecond): + } + close(release) + + // The worker is now inside the handler for "second", which means "first" is + // recorded and its sequence — the one catch-up started from — is satisfied. + // A snapshot-based catch-up returns here, with the newer publish unrecorded. + select { + case <-secondEntered: + case <-time.After(3 * time.Second): + t.Fatal("worker never reached the follow-up publish") + } + select { + case <-done: + t.Fatalf("catch-up returned on a stale target; diagnostics = %#v", sess.diagnosticsFor(uri)) + case <-time.After(25 * time.Millisecond): + } + + close(releaseSecond) + select { + case <-done: + case <-time.After(5 * time.Second): + t.Fatal("catch-up did not finish") + } + if got := sess.diagnosticsFor(uri); len(got) != 1 || got[0].Message != "second" { + t.Fatalf("diagnostics = %#v, want the publish accepted during catch-up", got) + } +} + +// TestCatchUpNotificationsGivesUpAfterTheLiveGrace spells out the tradeoff the +// live-client timeout represents: an unbounded wait would hang Check past a +// deadline the caller has already blown, so a worker that never catches up +// costs the grace and no more. Production sessions use a 50ms grace; this uses +// the same order of magnitude rather than the generous one the other tests pin. +func TestCatchUpNotificationsGivesUpAfterTheLiveGrace(t *testing.T) { + clientReader, serverWriter := io.Pipe() + defer serverWriter.Close() + client := NewClient(clientReader, io.Discard) + sess := &session{ + client: client, + catchUpGrace: 25 * time.Millisecond, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + blocked := make(chan struct{}) + release := make(chan struct{}) + defer close(release) + client.SetNotificationHandler(func(method string, _ json.RawMessage, _ int64) { + if method == "test/block" { + close(blocked) + <-release + } + }) + if err := writeMessage(serverWriter, map[string]any{"jsonrpc": "2.0", "method": "test/block"}); err != nil { + t.Fatal(err) + } + <-blocked + + done := make(chan struct{}) + start := time.Now() + go func() { + sess.catchUpNotifications() + close(done) + }() + select { + case <-done: + if elapsed := time.Since(start); elapsed < sess.catchUpGrace { + t.Fatalf("catch-up returned after %v, before its %v grace", elapsed, sess.catchUpGrace) + } + case <-time.After(3 * time.Second): + t.Fatal("catch-up on a permanently-behind worker must give up after the grace, not hang") + } +} + +func TestWaitForDiagnosticsCatchesUpOnContextCancellation(t *testing.T) { + clientReader, serverWriter := io.Pipe() + defer serverWriter.Close() + client := NewClient(clientReader, io.Discard) + uri := PathToURI("/repo/main.go") + sess := &session{ + client: client, + // Generous, so the catch-up is decided by the worker rather than by how + // fast this test's goroutines happen to be scheduled. + catchUpGrace: 5 * time.Second, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + blocked := make(chan struct{}) + release := make(chan struct{}) + client.SetNotificationHandler(func(method string, params json.RawMessage, seq int64) { + if method == "test/block" { + close(blocked) + <-release + return + } + sess.handleNotification(method, params, seq) + }) + + // Occupy the sole worker, then let a publish queue up behind it. + if err := writeMessage(serverWriter, map[string]any{"jsonrpc": "2.0", "method": "test/block"}); err != nil { + t.Fatal(err) + } + <-blocked + params := PublishDiagnosticsParams{URI: uri, Diagnostics: []Diagnostic{{Message: "fresh"}}} + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", "method": "textDocument/publishDiagnostics", "params": params, + }); err != nil { + t.Fatal(err) + } + deadline := time.Now().Add(time.Second) + for client.ReceiptSeq() != 2 { + if time.Now().After(deadline) { + t.Fatal("publish was not accepted before cancellation") + } + runtime.Gosched() + } + + ctx, cancel := context.WithCancel(context.Background()) + waitDone := make(chan bool, 1) + go func() { waitDone <- sess.waitForDiagnostics(ctx, uri, 0, 0) }() + deadline = time.Now().Add(time.Second) + for { + sess.mu.Lock() + waiting := len(sess.waiters[uri]) == 1 + sess.mu.Unlock() + if waiting { + break + } + if time.Now().After(deadline) { + t.Fatal("diagnostic wait was not registered") + } + runtime.Gosched() + } + + cancel() + select { + case fresh := <-waitDone: + t.Fatalf("wait returned %v while an accepted publish was still queued", fresh) + case <-time.After(25 * time.Millisecond): + } + close(release) + select { + case fresh := <-waitDone: + if !fresh { + t.Fatal("cancellation discarded a publish that was already off the wire") + } + case <-time.After(5 * time.Second): + t.Fatal("wait did not finish after the worker caught up") + } + if got := sess.diagnosticsFor(uri); len(got) != 1 || got[0].Message != "fresh" { + t.Fatalf("diagnostics = %#v, want the publish accepted before cancellation", got) + } +} + +// TestWaitForDiagnosticsDrainsFollowUpPublishOnCloseDuringDebounce covers +// jatmn's P2 finding (2): the debounce-phase closure arm returned true without +// draining, so Manager.Check read diagnosticsFor while a NEWER publish was still +// queued and handed back the older one. It also pins the drain as +// ctx-independent — cancelling mid-drain (finding (1)'s second clause) must not +// cut it short, since after closure the backlog can no longer grow. +func TestWaitForDiagnosticsDrainsFollowUpPublishOnCloseDuringDebounce(t *testing.T) { + clientReader, serverWriter := io.Pipe() + client := NewClient(clientReader, io.Discard) + uri := PathToURI("/repo/main.go") + sess := &session{ + client: client, + // Zero on purpose: the post-closure drain must not be bounded by the + // live-client grace. + catchUpGrace: 0, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + blocked := make(chan struct{}) + release := make(chan struct{}) + client.SetNotificationHandler(func(method string, params json.RawMessage, seq int64) { + if method == "test/block" { + close(blocked) + <-release + return + } + sess.handleNotification(method, params, seq) + }) + + publish := func(message string) { + t.Helper() + params := PublishDiagnosticsParams{URI: uri, Diagnostics: []Diagnostic{{Message: message}}} + if err := writeMessage(serverWriter, map[string]any{ + "jsonrpc": "2.0", "method": "textDocument/publishDiagnostics", "params": params, + }); err != nil { + t.Fatal(err) + } + } + + // A first publish lands and is handled, so the wait below starts out in its + // debounce phase rather than waiting for a first result. + publish("first") + deadline := time.Now().Add(time.Second) + for { + if got := sess.diagnosticsFor(uri); len(got) == 1 && got[0].Message == "first" { + break + } + if time.Now().After(deadline) { + t.Fatal("first publish was never handled") + } + runtime.Gosched() + } + // Now occupy the worker and queue a newer publish behind it. + if err := writeMessage(serverWriter, map[string]any{"jsonrpc": "2.0", "method": "test/block"}); err != nil { + t.Fatal(err) + } + <-blocked + publish("second") + deadline = time.Now().Add(time.Second) + for client.ReceiptSeq() != 3 { + if time.Now().After(deadline) { + t.Fatal("follow-up publish was not accepted before closure") + } + runtime.Gosched() + } + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + waitDone := make(chan bool, 1) + go func() { waitDone <- sess.waitForDiagnostics(ctx, uri, 5*time.Second, 0) }() + select { + case fresh := <-waitDone: + t.Fatalf("wait returned %v instead of debouncing", fresh) + case <-time.After(25 * time.Millisecond): + } + + if err := serverWriter.Close(); err != nil { + t.Fatal(err) + } + deadline = time.Now().Add(time.Second) + for !client.IsClosed() { + if time.Now().After(deadline) { + t.Fatal("transport EOF did not close client") + } + runtime.Gosched() + } + // Cancelling now must not cut the drain short. + cancel() + select { + case fresh := <-waitDone: + t.Fatalf("wait returned %v before the accepted follow-up publish drained", fresh) + case <-time.After(25 * time.Millisecond): + } + + close(release) + select { + case fresh := <-waitDone: + if !fresh { + t.Fatal("wait dropped the publish it had already observed") + } + case <-time.After(time.Second): + t.Fatal("wait did not finish after the notification drain") + } + if got := sess.diagnosticsFor(uri); len(got) != 1 || got[0].Message != "second" { + t.Fatalf("diagnostics = %#v, want the follow-up publish accepted before closure", got) + } +} + +func TestWaitForDiagnosticsPreservesPublishHandledAtDeadline(t *testing.T) { + previousProcs := runtime.GOMAXPROCS(1) + defer runtime.GOMAXPROCS(previousProcs) + params, err := json.Marshal(PublishDiagnosticsParams{ + URI: PathToURI("/repo/main.go"), Diagnostics: []Diagnostic{{Message: "fresh"}}, + }) + if err != nil { + t.Fatal(err) + } + for i := 0; i < 100; i++ { + client := &Client{closed: make(chan struct{})} + uri := PathToURI("/repo/main.go") + sess := &session{ + client: client, + diagnostics: map[string][]Diagnostic{}, + lastPublish: map[string]time.Time{}, + publishSeq: map[string]int64{}, + waiters: map[string][]chan struct{}{}, + } + ctx, cancel := context.WithCancel(context.Background()) + waitDone := make(chan bool, 1) + go func() { + waitDone <- sess.waitForDiagnostics(ctx, uri, time.Second, 0) + }() + + deadline := time.Now().Add(time.Second) + for { + sess.mu.Lock() + waiting := len(sess.waiters[uri]) == 1 + sess.mu.Unlock() + if waiting { + break + } + if time.Now().After(deadline) { + t.Fatal("diagnostic wait was not registered") + } + runtime.Gosched() + } + + // handleNotification records the publish and closes the waiter. Cancel + // before yielding under GOMAXPROCS(1), making both cases ready together. + sess.handleNotification("textDocument/publishDiagnostics", params, 1) + cancel() + + select { + case fresh := <-waitDone: + if !fresh { + t.Fatalf("iteration %d: deadline discarded a publish handled before cancellation", i) + } + case <-time.After(time.Second): + t.Fatalf("iteration %d: wait did not return", i) + } + } +} + func TestManagerCheckDegradesWhenServerBinaryMissing(t *testing.T) { // A configured extension whose binary isn't on PATH (exec.ErrNotFound) must // degrade to no diagnostics, exactly like an unsupported extension.