perf: reuse one Docutils interpreter instead of one per file

Signed-off-by: Joseph Kato <joseph@jdkato.io>

Author: jdkato

internal/lint/lint.go

@@ -41,6 +41,10 @@ type Linter struct {
 	mdx     *procPool
 	mdxOnce sync.Once
 
+	// rst is the same arrangement for Docutils.
+	rst     *procPool
+	rstOnce sync.Once
+
 	// inScope lists the rules whose scope matches a given block scope, keyed by
 	// the block's scope and parent.
 	//
@@ -465,4 +469,9 @@ func (l *Linter) stopExternal() {
 		l.mdx = nil
 		l.mdxOnce = sync.Once{}
 	}
+	if l.rst != nil {
+		l.rst.stop()
+		l.rst = nil
+		l.rstOnce = sync.Once{}
+	}
 }

internal/lint/rst.go

@@ -1,9 +1,13 @@
 package lint
 
 import (
+	"bufio"
 	"errors"
+	"os"
+	"path/filepath"
 	"regexp"
 	"strings"
+	"sync"
 
 	"github.com/errata-ai/vale/v3/internal/core"
 	"github.com/errata-ai/vale/v3/internal/system"
@@ -30,6 +34,125 @@ var rstArgs = []string{
 	"--no-section-numbering",
 }
 
+// Converting a document with Docutils takes a few milliseconds; starting
+// Python and importing Docutils takes seventy. Vale paid the latter once per
+// file, so on a corpus of reStructuredText nearly all of the time went to
+// starting the same interpreter over and over.
+//
+// rstServer keeps one interpreter up and converts documents as they arrive. It
+// parses the flags through Docutils' own command-line handling rather than
+// mapping them to settings by hand, so the pooled settings cannot drift from
+// what a per-file rst2html invocation would use.
+const rstServer = `import sys
+from docutils.core import Publisher, publish_string
+
+_pub = Publisher()
+_pub.set_components("standalone", "restructuredtext", "html4css1")
+_pub.process_command_line(argv=sys.argv[1:])
+SETTINGS = _pub.settings
+
+buf = sys.stdin.buffer
+out = sys.stdout.buffer
+
+while True:
+    header = buf.readline()
+    if not header:
+        break
+    n = int(header)
+    if n < 0:
+        break
+    doc = buf.read(n) if n else b""
+    try:
+        b = publish_string(
+            source=doc.decode("utf-8"),
+            source_path="<stdin>",
+            writer_name="html4css1",
+            settings=SETTINGS,
+        )
+        out.write(b"ok %d\n" % len(b))
+        out.write(b)
+    except Exception as e:
+        m = str(e).encode("utf-8", "replace")
+        out.write(b"err %d\n" % len(m))
+        out.write(m)
+    out.flush()`
+
+var (
+	rstOnce   sync.Once
+	rstDirect []string // <python> -c <server>, or nil
+)
+
+// rstInterpreter finds the Python that can import Docutils.
+//
+// It is not necessarily the `python3` on PATH: rst2html is installed with a
+// shebang naming the interpreter of the environment Docutils was installed
+// into, and on a machine with several Pythons the one on PATH often cannot
+// import it. So the script is asked which interpreter it runs on.
+func rstInterpreter(exe string) string {
+	resolved, err := filepath.EvalSymlinks(exe)
+	if err != nil {
+		return ""
+	}
+
+	file, err := os.Open(resolved)
+	if err != nil {
+		return ""
+	}
+	defer file.Close()
+
+	line, err := bufio.NewReader(file).ReadString('\n')
+	if err != nil || !strings.HasPrefix(line, "#!") {
+		return ""
+	}
+
+	fields := strings.Fields(strings.TrimPrefix(strings.TrimSpace(line), "#!"))
+	if len(fields) == 0 {
+		return ""
+	}
+
+	// `#!/usr/bin/env python3` names the interpreter in the second field.
+	if filepath.Base(fields[0]) == "env" {
+		if len(fields) < 2 {
+			return ""
+		}
+		return system.Which([]string{fields[1]})
+	}
+
+	return fields[0]
+}
+
+// rstFastPath returns the argv prefix for converting through a long-lived
+// interpreter, or nil when that could not be established.
+func rstFastPath(exe string) []string {
+	rstOnce.Do(func() {
+		python := rstInterpreter(exe)
+		if python == "" {
+			return
+		}
+
+		candidate := []string{python, "-c", rstServer}
+
+		// Trust it only after a document has made the round trip. The probe
+		// carries a non-ASCII character on purpose: a mismatched default
+		// encoding is what broke the first version of the AsciiDoc pool, and
+		// an ASCII-only probe would have passed anyway.
+		probe, err := startExtProc(candidate, rstArgs)
+		if err != nil {
+			return
+		}
+		defer probe.close()
+
+		got, err := probe.convert("naïve body\n")
+		if err != nil || !strings.Contains(got, "naïve body") {
+			return
+		}
+
+		rstDirect = candidate
+	})
+
+	return rstDirect
+}
+
 func (l *Linter) lintRST(f *core.File) error {
 	var html string
 
@@ -53,19 +176,44 @@ func (l *Linter) lintRST(f *core.File) error {
 	s = reSphinx.ReplaceAllString(s, ".. code::")
 	s = reCodeBlock.ReplaceAllString(s, "::")
 
-	html, err = callRst(s, rst2html)
+	html, err = l.callRst(s, rst2html)
 	if err != nil {
 		return core.NewE100(f.Path, err)
 	}
 
 	return l.lintHTMLTokens(f, []byte(html), 0)
 }
 
-func callRst(text, lib string) (string, error) {
-	html, err := system.ExecuteWithInput(lib, text, rstArgs...)
+// callRst converts one document, over a pooled interpreter when Docutils can
+// be reached directly.
+func (l *Linter) callRst(text, exe string) (string, error) {
+	if direct := rstFastPath(exe); direct != nil {
+		l.rstOnce.Do(func() {
+			pool, err := newProcPool(direct, rstArgs, adocConcurrency)
+			if err == nil {
+				l.rst = pool
+			}
+		})
+
+		if l.rst != nil {
+			html, err := l.rst.convert(text, direct, rstArgs)
+			if err != nil {
+				return "", err
+			}
+			return rstBody(html), nil
+		}
+	}
+
+	html, err := system.ExecuteWithInput(exe, text, rstArgs...)
 	if err != nil {
 		return "", err
 	}
+
+	return rstBody(html), nil
+}
+
+// rstBody takes the document body out of a full rst2html page.
+func rstBody(html string) string {
 	html = strings.ReplaceAll(html, "\r", "")
 
 	bodyStart := strings.Index(html, "<body>\n")
@@ -80,5 +228,5 @@ func callRst(text, lib string) (string, error) {
 		}
 	}
 
-	return html[bodyStart+7 : bodyEnd], nil
+	return html[bodyStart+7 : bodyEnd]
 }

internal/lint/rst_pool_test.go

@@ -0,0 +1,148 @@
+package lint
+
+import (
+	"os/exec"
+	"runtime"
+	"strings"
+	"testing"
+
+	"github.com/errata-ai/vale/v3/internal/system"
+)
+
+const rstFixture = "../../testdata/fixtures/formats/test.rst"
+
+func rstExecutable() string {
+	return system.Which([]string{
+		"rst2html", "rst2html.py", "rst2html-3", "rst2html-3.py"})
+}
+
+// The pool reuses one interpreter for every document, so its Docutils settings
+// are established once instead of per file. That is the whole risk of the
+// change: if the pooled settings differ at all from what a per-file rst2html
+// invocation uses, Vale silently lints different HTML.
+func TestRSTPoolMatchesSpawnedOutput(t *testing.T) {
+	exe := rstExecutable()
+	if exe == "" {
+		t.Skip("rst2html not installed")
+	}
+
+	argv := rstFastPath(exe)
+	if argv == nil {
+		t.Skip("docutils not reachable from the rst2html interpreter")
+	}
+
+	pool, err := newProcPool(argv, rstArgs, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer pool.stop()
+
+	// Non-ASCII and a title on purpose: a mismatched default encoding and the
+	// source name reaching <title> are both ways the two paths can diverge
+	// without any visible error.
+	for _, doc := range []string{
+		"Title\n=====\n\nnaïve body — with punctuation.\n",
+		"no title here, just a paragraph\n",
+		"- a list\n- of items\n\n``literal`` and *emphasis*\n",
+		// Section numbering, table-of-contents backlinks and footnote
+		// backlinks are each turned off by a flag in rstArgs, and each shows
+		// up inside <body>. Without a document like this the comparison
+		// passes even if the pooled interpreter ignores the flags entirely.
+		".. sectnum::\n\n.. contents::\n\nFirst\n=====\n\n" +
+			"Body with a footnote [1]_.\n\nSecond\n======\n\nMore.\n\n.. [1] The note.\n",
+	} {
+		pooled, cErr := pool.convert(doc, argv, rstArgs)
+		if cErr != nil {
+			t.Fatalf("pooled convert failed: %v", cErr)
+		}
+
+		spawned, sErr := system.ExecuteWithInput(exe, doc, rstArgs...)
+		if sErr != nil {
+			t.Fatalf("spawned convert failed: %v", sErr)
+		}
+
+		if got, want := rstBody(pooled), rstBody(spawned); got != want {
+			t.Errorf("pooled and spawned output differ for %q:\n pooled: %q\nspawned: %q",
+				doc, got, want)
+		}
+	}
+}
+
+// A pool with no owner is a pool nothing stops. The CLI gets away with it --
+// exiting reaps the children -- but the language server lints many times in one
+// process, and every run would leave its interpreters behind.
+func TestRSTPoolDoesNotLeakProcesses(t *testing.T) {
+	if exe := rstExecutable(); exe == "" || rstFastPath(exe) == nil {
+		t.Skip("docutils not reachable on this machine")
+	}
+
+	before := pythonProcessCount(t)
+
+	linter, err := initLinter()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// A real reStructuredText file, so the run actually starts the pool -- a
+	// path that lints nothing would pass this test without testing anything.
+	linted, err := linter.Lint([]string{rstFixture}, "*")
+	if err != nil {
+		t.Fatal(err)
+	}
+	if len(linted) == 0 {
+		t.Fatal("nothing linted; the pool was never started")
+	}
+	if linter.rst != nil {
+		t.Error("pool outlived the run")
+	}
+
+	// Give the OS a moment to reap.
+	runtime.Gosched()
+	if after := pythonProcessCount(t); after > before {
+		t.Errorf("leaked processes: %d before, %d after", before, after)
+	}
+}
+
+// The pool has to survive a document Docutils rejects: later files still
+// convert, rather than the run losing its warm interpreters.
+func TestRSTPoolSurvivesABadDocument(t *testing.T) {
+	exe := rstExecutable()
+	if exe == "" {
+		t.Skip("rst2html not installed")
+	}
+
+	argv := rstFastPath(exe)
+	if argv == nil {
+		t.Skip("docutils not reachable from the rst2html interpreter")
+	}
+
+	pool, err := newProcPool(argv, rstArgs, 1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer pool.stop()
+
+	// A malformed directive, which Docutils reports at the severity that
+	// --halt=5 stops on.
+	if _, err = pool.convert(".. |bad\xff| replace::\n", argv, rstArgs); err == nil {
+		t.Log("bad document was accepted; the recovery path is untested here")
+	}
+
+	got, err := pool.convert("still working\n", argv, rstArgs)
+	if err != nil {
+		t.Fatalf("pool did not recover: %v", err)
+	}
+	if !strings.Contains(got, "still working") {
+		t.Errorf("unexpected output after recovery: %q", got)
+	}
+}
+
+func pythonProcessCount(t *testing.T) int {
+	t.Helper()
+
+	out, err := exec.Command("ps", "-eo", "comm").Output()
+	if err != nil {
+		t.Skip("ps unavailable")
+	}
+	return strings.Count(string(out), "python")
+}