Autocomplete: wrap some parser.parse() calls in OpenTelemetry spans (…

…#3419)

agent/src/cli/evaluate-autocomplete/AutocompleteMatcher.ts

@@ -1,8 +1,8 @@
 import * as vscode from 'vscode'
-import type { Tree, default as Parser } from 'web-tree-sitter'
+import type { Tree } from 'web-tree-sitter'
 
 import { SupportedLanguage, isSupportedLanguage } from '../../../../vscode/src/tree-sitter/grammars'
-import { createParser } from '../../../../vscode/src/tree-sitter/parser'
+import { type WrappedParser, createParser } from '../../../../vscode/src/tree-sitter/parser'
 
 import { EvaluationDocument, type EvaluationDocumentParams } from './EvaluationDocument'
 import type { Queries } from './Queries'
@@ -20,7 +20,7 @@ interface AutocompleteMatch {
     requestPosition: vscode.Position
 }
 export class AutocompleteMatcher {
-    public parser: Parser | undefined
+    public parser: WrappedParser | undefined
     public originalTree: Tree | undefined
     public originalTreeIsFreeOfErrrors: boolean | undefined
     constructor(

agent/src/cli/evaluate-autocomplete/Queries.ts

@@ -1,9 +1,10 @@
 import * as path from 'path'
 import * as fspromises from 'fs/promises'
 
-import type { Query, default as Parser } from 'web-tree-sitter'
+import type { Query } from 'web-tree-sitter'
 
 import { SupportedLanguage } from '../../../../vscode/src/tree-sitter/grammars'
+import type { WrappedParser } from '../../../../vscode/src/tree-sitter/parser'
 
 type QueryName = 'context'
 
@@ -19,7 +20,7 @@ export class Queries {
     private cache: CompiledQuery[] = []
     constructor(private queriesDirectory: string) {}
     public async loadQuery(
-        parser: Parser,
+        parser: WrappedParser,
         language: SupportedLanguage,
         name: QueryName
     ): Promise<Query | undefined> {
@@ -38,7 +39,7 @@ export class Queries {
     }
 
     private async compileQuery(
-        parser: Parser,
+        parser: WrappedParser,
         language: SupportedLanguage,
         name: QueryName
     ): Promise<Query | undefined> {
@@ -93,7 +94,7 @@ interface CompiledQuery extends UncompiledQuery {
     compiledQuery: Query
 }
 
-function compileQuery(query: UncompiledQuery, parser: Parser): CompiledQuery {
+function compileQuery(query: UncompiledQuery, parser: WrappedParser): CompiledQuery {
     return {
         ...query,
         compiledQuery: parser.getLanguage().query(query.queryString),

agent/src/cli/evaluate-autocomplete/testParse.ts

@@ -1,7 +1,7 @@
-import type Parser from 'web-tree-sitter'
+import type { WrappedParser } from '../../../../vscode/src/tree-sitter/parser'
 
 /** Returns true if the new text parses successfully. */
-export function testParses(newText: string, parser: Parser): boolean | undefined {
+export function testParses(newText: string, parser: WrappedParser): boolean | undefined {
     // Originally, this function passed the `previousTree` argument to benefit
     // from performance improvements but it didn't work correctly,
     // parseTest.test.ts was failing until we removed `previousTree`.

agent/src/cli/evaluate-autocomplete/triggerAutocomplete.ts

@@ -1,12 +1,13 @@
 import { calcPatch } from 'fast-myers-diff'
 import * as vscode from 'vscode'
-import type { Tree, default as Parser } from 'web-tree-sitter'
+import type { Tree } from 'web-tree-sitter'
 
 import { ProtocolTextDocumentWithUri } from '../../../../vscode/src/jsonrpc/TextDocumentWithUri'
 import { AgentTextDocument } from '../../AgentTextDocument'
 import type { MessageHandler } from '../../jsonrpc-alias'
 import type { AutocompleteResult } from '../../protocol-alias'
 
+import type { WrappedParser } from '../../../../vscode/src/tree-sitter/parser'
 import type { AutocompleteMatchKind } from './AutocompleteMatcher'
 import type { EvaluationDocument } from './EvaluationDocument'
 import type { TestParameters } from './TestParameters'
@@ -15,7 +16,7 @@ import { testParses } from './testParse'
 import { testTypecheck } from './testTypecheck'
 
 export interface AutocompleteParameters {
-    parser?: Parser
+    parser?: WrappedParser
     originalTree?: Tree
     originalTreeIsErrorFree?: boolean
     client: MessageHandler

vscode/src/completions/get-current-doc-context.test.ts

@@ -5,7 +5,7 @@ import type * as Parser from 'web-tree-sitter'
 
 import { range } from '../testutils/textDocument'
 import { asPoint } from '../tree-sitter/parse-tree-cache'
-import { resetParsersCache } from '../tree-sitter/parser'
+import { type WrappedParser, resetParsersCache } from '../tree-sitter/parser'
 
 import { getContextRange } from './doc-context-getters'
 import {
@@ -194,7 +194,7 @@ describe('getCurrentDocContext', () => {
     })
 
     describe('multiline triggers', () => {
-        let parser: Parser
+        let parser: WrappedParser
 
         interface PrepareTestParams {
             code: string

vscode/src/completions/text-processing/parse-completion.ts

@@ -6,6 +6,7 @@ import { asPoint, getCachedParseTreeForDocument } from '../../tree-sitter/parse-
 import type { DocumentContext } from '../get-current-doc-context'
 import type { InlineCompletionItem } from '../types'
 
+import type { WrappedParser } from '../../tree-sitter/parser'
 import {
     type InlineCompletionItemWithAnalytics,
     getMatchingSuffixLength,
@@ -89,7 +90,7 @@ interface PasteCompletionParams {
     document: TextDocument
     docContext: DocumentContext
     tree: Tree
-    parser: Parser
+    parser: WrappedParser
 }
 
 interface PasteCompletionResult {
@@ -129,7 +130,7 @@ function pasteCompletion(params: PasteCompletionParams): PasteCompletionResult {
 
     // TODO(tree-sitter): consider parsing only the changed part of the document to improve performance.
     // parser.parse(textWithCompletion, tree, { includedRanges: [...]})
-    const treeWithCompletion = parser.parse(textWithCompletion, treeCopy)
+    const treeWithCompletion = parser.observableParse(textWithCompletion, treeCopy)
     addAutocompleteDebugEvent('paste-completion', {
         text: textWithCompletion,
     })

vscode/src/tree-sitter/parse-tree-cache.ts

@@ -4,15 +4,15 @@ import type { TextDocument } from 'vscode'
 import type { Tree, default as Parser } from 'web-tree-sitter'
 
 import { type SupportedLanguage, isSupportedLanguage } from './grammars'
-import { createParser, getParser } from './parser'
+import { type WrappedParser, createParser, getParser } from './parser'
 
 const parseTreesPerFile = new LRUCache<string, Tree>({
     max: 10,
 })
 
 interface ParseTreeCache {
     tree: Tree
-    parser: Parser
+    parser: WrappedParser
     cacheKey: string
 }
 
@@ -49,7 +49,7 @@ async function parseDocument(document: TextDocument): Promise<void> {
     updateParseTreeCache(document, parser)
 }
 
-export function updateParseTreeCache(document: TextDocument, parser: Parser): void {
+export function updateParseTreeCache(document: TextDocument, parser: WrappedParser): void {
     const tree = parser.parse(document.getText())
     parseTreesPerFile.set(document.uri.toString(), tree)
 }

vscode/src/tree-sitter/parser.ts

@@ -3,6 +3,7 @@ import path from 'path'
 import * as vscode from 'vscode'
 import type Parser from 'web-tree-sitter'
 
+import { wrapInActiveSpan } from '@sourcegraph/cody-shared'
 import type { Tree } from 'web-tree-sitter'
 import { DOCUMENT_LANGUAGE_TO_GRAMMAR, type SupportedLanguage, isSupportedLanguage } from './grammars'
 import { initQueries } from './query-sdk'
@@ -14,7 +15,7 @@ const ParserImpl = require('web-tree-sitter') as typeof Parser
  * and load language grammar only once, first time we need parser for a specific
  * language, next time we read it from this cache.
  */
-const PARSERS_LOCAL_CACHE: Partial<Record<SupportedLanguage, Parser>> = {}
+const PARSERS_LOCAL_CACHE: Partial<Record<SupportedLanguage, WrappedParser>> = {}
 
 interface ParserSettings {
     language: SupportedLanguage
@@ -26,7 +27,7 @@ interface ParserSettings {
     grammarDirectory?: string
 }
 
-export function getParser(language: SupportedLanguage): Parser | undefined {
+export function getParser(language: SupportedLanguage): WrappedParser | undefined {
     return PARSERS_LOCAL_CACHE[language]
 }
 
@@ -45,7 +46,14 @@ async function isRegularFile(uri: vscode.Uri): Promise<boolean> {
     }
 }
 
-export async function createParser(settings: ParserSettings): Promise<Parser | undefined> {
+export type WrappedParser = Pick<Parser, 'parse' | 'getLanguage'> & {
+    /**
+     * Wraps `parser.parse()` call into an OpenTelemetry span.
+     */
+    observableParse: Parser['parse']
+}
+
+export async function createParser(settings: ParserSettings): Promise<WrappedParser | undefined> {
     const { language, grammarDirectory = __dirname } = settings
 
     const cachedParser = PARSERS_LOCAL_CACHE[language]
@@ -66,14 +74,24 @@ export async function createParser(settings: ParserSettings): Promise<Parser | u
     const languageGrammar = await ParserImpl.Language.load(wasmPath)
 
     parser.setLanguage(languageGrammar)
-    // stop parsing after 70ms to avoid infinite loops
-    // if that happens, tree-sitter throws an error so we can catch and address it
-    parser.setTimeoutMicros(70_000)
-    PARSERS_LOCAL_CACHE[language] = parser
 
+    // Disable the timeout in unit tests to avoid timeout errors.
+    if (!process.env.VITEST) {
+        // Stop parsing after 70ms to avoid infinite loops.
+        // If that happens, tree-sitter throws an error so we can catch and address it.
+        parser.setTimeoutMicros(70_000)
+    }
+
+    const wrappedParser: WrappedParser = {
+        getLanguage: () => parser.getLanguage(),
+        parse: (...args) => parser.parse(...args),
+        observableParse: (...args) => wrapInActiveSpan('parser.parse', () => parser.parse(...args)),
+    }
+
+    PARSERS_LOCAL_CACHE[language] = wrappedParser
     initQueries(languageGrammar, language, parser)
 
-    return parser
+    return wrappedParser
 }
 
 export function parseString(languageId: string, source: string): Tree | null {

vscode/src/tree-sitter/query-sdk.ts

@@ -12,7 +12,7 @@ import type {
 
 import { type SupportedLanguage, isSupportedLanguage } from './grammars'
 import { getCachedParseTreeForDocument } from './parse-tree-cache'
-import { getParser } from './parser'
+import { type WrappedParser, getParser } from './parser'
 import { type CompletionIntent, type QueryName, intentPriority, languages } from './queries'
 
 interface ParsedQuery {
@@ -59,7 +59,7 @@ export function initQueries(language: Language, languageId: SupportedLanguage, p
 }
 
 export interface DocumentQuerySDK {
-    parser: Parser
+    parser: WrappedParser
     queries: ResolvedQueries & QueryWrappers
     language: SupportedLanguage
 }

vscode/src/tree-sitter/query-tests/annotate-and-match-snapshot.ts

@@ -8,6 +8,7 @@ import type { Point, SyntaxNode, default as Parser } from 'web-tree-sitter'
 
 import type { SupportedLanguage } from '../grammars'
 import { getLanguageConfig } from '../language'
+import type { WrappedParser } from '../parser'
 
 interface CommentSymbolInfo {
     delimiter: string
@@ -104,7 +105,7 @@ export type Captures = (
 interface AnnotateSnippetParams {
     code: string
     language: SupportedLanguage
-    parser: Parser
+    parser: WrappedParser
     captures: Captures
     isOnly: boolean
 }
@@ -235,7 +236,7 @@ function commentOutLines(text: string, commentSymbol: string): string {
 
 interface AnnotateAndMatchParams {
     sourcesPath: string
-    parser: Parser
+    parser: WrappedParser
     language: SupportedLanguage
     captures: Captures
 }

vscode/src/tree-sitter/test-helpers.ts

@@ -1,9 +1,9 @@
 import path from 'path'
 
-import type { QueryCapture, QueryMatch, default as Parser } from 'web-tree-sitter'
+import type { QueryCapture, QueryMatch } from 'web-tree-sitter'
 
 import { SupportedLanguage } from './grammars'
-import { createParser } from './parser'
+import { type WrappedParser, createParser } from './parser'
 import { type DocumentQuerySDK, getDocumentQuerySDK } from './query-sdk'
 
 const CUSTOM_WASM_LANGUAGE_DIR = path.join(__dirname, '../../resources/wasm')
@@ -13,7 +13,7 @@ const CUSTOM_WASM_LANGUAGE_DIR = path.join(__dirname, '../../resources/wasm')
  */
 export function initTreeSitterParser(
     language = SupportedLanguage.typescript
-): Promise<Parser | undefined> {
+): Promise<WrappedParser | undefined> {
     return createParser({
         language,
         grammarDirectory: CUSTOM_WASM_LANGUAGE_DIR,