Bound document content extraction time and decouple it from uploads (MM-69098) (#36856) (#37043)

Automatic Merge

server/platform/services/docextractor/docextractor.go

@@ -4,7 +4,9 @@
 package docextractor
 
 import (
+	"fmt"
 	"io"
+	"time"
 
 	"github.com/mattermost/mattermost/server/public/shared/mlog"
 )
@@ -15,6 +17,23 @@ type ExtractSettings struct {
 	MaxFileSize      int64
 	MMPreviewURL     string
 	MMPreviewSecret  string
+	// Timeout bounds how long a caller waits for a single extraction. A value
+	// <= 0 disables it. NOTE: this bounds wall-clock wait time (and thus how
+	// long an extraction occupies its caller's worker slot), NOT CPU work.
+	// The docconv converters are not context-aware, so on timeout the
+	// converter keeps running to completion on a detached goroutine and keeps
+	// consuming CPU until it finishes on its own. Under sustained load,
+	// detached extractions can therefore accumulate and run concurrently. The
+	// primary bound on the work of any single extraction is MaxFileSize, which
+	// limits how much input the converter reads.
+	Timeout time.Duration
+	// ReaderCloser, when set, transfers ownership of closing the input reader
+	// to this package. It is closed only after extraction has actually
+	// finished reading. This matters with Timeout set: on timeout the caller
+	// returns while the converter may still be reading on a detached
+	// goroutine, so the caller must NOT close the reader itself or it would
+	// race with (and close the file out from under) that goroutine.
+	ReaderCloser io.Closer
 }
 
 // Extract extract the text from a document using the system default extractors
@@ -45,7 +64,71 @@ func ExtractWithExtraExtractors(logger mlog.LoggerIFace, filename string, r io.R
 	enabledExtractors.Add(&plainExtractor{})
 
 	if enabledExtractors.Match(filename) {
-		return enabledExtractors.Extract(filename, r, settings.MaxFileSize)
+		return extractWithTimeout(enabledExtractors, filename, r, settings)
+	}
+
+	// No extractor matched, so nothing will read r; close it here since
+	// extractWithTimeout (which otherwise owns the close) is never reached.
+	if settings.ReaderCloser != nil {
+		settings.ReaderCloser.Close()
 	}
 	return "", nil
 }
+
+// extractWithTimeout runs the extraction and stops waiting for it once
+// settings.Timeout elapses. Because the underlying docconv converters are not
+// context-aware, the extraction runs on a detached goroutine: on timeout we
+// stop waiting and return an error, releasing the caller (and its worker slot)
+// even though the converter keeps running.
+//
+// This decouples extraction from the caller, but it does NOT cap CPU: the
+// detached converter continues to completion in the background, so a sustained
+// stream of expensive documents can leave several detached extractions running
+// at once. The per-extraction work is bounded instead by MaxFileSize (input
+// size). Load-shedding on the number of in-flight detached extractions is a
+// possible future improvement; it is intentionally not done here so it does
+// not also throttle the backfill job that re-extracts skipped content.
+func extractWithTimeout(e Extractor, filename string, r io.ReadSeeker, settings ExtractSettings) (string, error) {
+	if settings.Timeout <= 0 {
+		if settings.ReaderCloser != nil {
+			defer settings.ReaderCloser.Close()
+		}
+		return e.Extract(filename, r, settings.MaxFileSize)
+	}
+
+	type extractResult struct {
+		text string
+		err  error
+	}
+	resultCh := make(chan extractResult, 1)
+	go func() {
+		// This goroutine owns the reader for the lifetime of the extraction.
+		// After the timeout fires the caller returns, but the converter may
+		// still be reading r here, so the reader is closed only once this
+		// goroutine is done with it - never by the caller.
+		if settings.ReaderCloser != nil {
+			defer settings.ReaderCloser.Close()
+		}
+		// This goroutine is detached, so an unrecovered panic in an extractor
+		// would crash the whole server. Convert it into an error instead.
+		// resultCh is buffered (cap 1), so this send never blocks even if the
+		// caller already timed out and stopped receiving.
+		defer func() {
+			if rec := recover(); rec != nil {
+				resultCh <- extractResult{err: fmt.Errorf("panic during document text extraction: %v", rec)}
+			}
+		}()
+		text, err := e.Extract(filename, r, settings.MaxFileSize)
+		resultCh <- extractResult{text: text, err: err}
+	}()
+
+	timer := time.NewTimer(settings.Timeout)
+	defer timer.Stop()
+
+	select {
+	case res := <-resultCh:
+		return res.text, res.err
+	case <-timer.C:
+		return "", fmt.Errorf("document text extraction timed out after %s", settings.Timeout)
+	}
+}

server/platform/services/docextractor/docextractor_test.go

@@ -8,7 +8,9 @@ import (
 	"errors"
 	"io"
 	"strings"
+	"sync/atomic"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -213,6 +215,144 @@ func TestExtractWithExtraExtractors(t *testing.T) {
 	})
 }
 
+type slowExtractor struct {
+	delay time.Duration
+}
+
+func (se *slowExtractor) Name() string { return "slowExtractor" }
+
+func (se *slowExtractor) Match(filename string) bool { return true }
+
+func (se *slowExtractor) Extract(filename string, r io.ReadSeeker, _ int64) (string, error) {
+	time.Sleep(se.delay)
+	return "done", nil
+}
+
+func TestExtractTimeout(t *testing.T) {
+	logger := mlog.CreateConsoleTestLogger(t)
+	data := []byte("hello world")
+
+	t.Run("aborts a slow extraction once the timeout elapses", func(t *testing.T) {
+		start := time.Now()
+		text, err := ExtractWithExtraExtractors(logger, "file.txt", bytes.NewReader(data), ExtractSettings{Timeout: 50 * time.Millisecond}, []Extractor{&slowExtractor{delay: 10 * time.Second}})
+		elapsed := time.Since(start)
+
+		require.Error(t, err)
+		require.Empty(t, text)
+		assert.Contains(t, err.Error(), "timed out")
+		assert.Less(t, elapsed, 5*time.Second, "should return shortly after the timeout, not wait for the extraction")
+	})
+
+	t.Run("returns the result when extraction finishes within the timeout", func(t *testing.T) {
+		text, err := ExtractWithExtraExtractors(logger, "file.txt", bytes.NewReader(data), ExtractSettings{Timeout: 5 * time.Second}, []Extractor{&slowExtractor{delay: 10 * time.Millisecond}})
+		require.NoError(t, err)
+		require.Equal(t, "done", text)
+	})
+
+	t.Run("a zero timeout disables the bound", func(t *testing.T) {
+		text, err := ExtractWithExtraExtractors(logger, "file.txt", bytes.NewReader(data), ExtractSettings{Timeout: 0}, []Extractor{&slowExtractor{delay: 10 * time.Millisecond}})
+		require.NoError(t, err)
+		require.Equal(t, "done", text)
+	})
+
+	t.Run("a panic in the detached extraction is converted to an error", func(t *testing.T) {
+		text, err := ExtractWithExtraExtractors(logger, "file.txt", bytes.NewReader(data), ExtractSettings{Timeout: time.Second}, []Extractor{&panickingExtractor{}})
+		require.Error(t, err)
+		require.Empty(t, text)
+		require.Contains(t, err.Error(), "panic")
+	})
+}
+
+type panickingExtractor struct{}
+
+func (pe *panickingExtractor) Name() string { return "panickingExtractor" }
+
+func (pe *panickingExtractor) Match(filename string) bool { return true }
+
+func (pe *panickingExtractor) Extract(filename string, r io.ReadSeeker, _ int64) (string, error) {
+	panic("boom")
+}
+
+type recordingCloser struct {
+	closed atomic.Bool
+}
+
+func (c *recordingCloser) Close() error {
+	c.closed.Store(true)
+	return nil
+}
+
+// blockingExtractor blocks inside Extract until release is closed, simulating a
+// converter that is still using the reader after an extraction timeout fires.
+type blockingExtractor struct {
+	started chan struct{}
+	release chan struct{}
+}
+
+func (be *blockingExtractor) Name() string { return "blockingExtractor" }
+
+func (be *blockingExtractor) Match(filename string) bool { return true }
+
+func (be *blockingExtractor) Extract(filename string, r io.ReadSeeker, _ int64) (string, error) {
+	close(be.started)
+	<-be.release
+	return "done", nil
+}
+
+func TestExtractReaderCloserOwnership(t *testing.T) {
+	logger := mlog.CreateConsoleTestLogger(t)
+
+	t.Run("reader is closed only after the detached extraction finishes on timeout", func(t *testing.T) {
+		closer := &recordingCloser{}
+		be := &blockingExtractor{started: make(chan struct{}), release: make(chan struct{})}
+		settings := ExtractSettings{Timeout: 50 * time.Millisecond, ReaderCloser: closer}
+
+		_, err := ExtractWithExtraExtractors(logger, "file.txt", bytes.NewReader([]byte("hi")), settings, []Extractor{be})
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "timed out")
+
+		// Wait (with a deadline) for the detached extraction to start so the
+		// test fails fast instead of hanging if it never runs.
+		select {
+		case <-be.started:
+		case <-time.After(2 * time.Second):
+			require.FailNow(t, "extraction did not start within the deadline")
+		}
+		// The extraction goroutine is still running, so closing the reader now
+		// would race with it; it must stay open.
+		require.False(t, closer.closed.Load(), "reader must not be closed while the extraction goroutine is still running")
+
+		close(be.release)
+		require.Eventually(t, closer.closed.Load, 2*time.Second, 5*time.Millisecond, "reader should be closed once the extraction goroutine finishes")
+	})
+
+	t.Run("reader is closed on the synchronous path", func(t *testing.T) {
+		closer := &recordingCloser{}
+		_, err := ExtractWithExtraExtractors(logger, "file.txt", bytes.NewReader([]byte("hi")), ExtractSettings{ReaderCloser: closer}, []Extractor{&slowExtractor{delay: 0}})
+		require.NoError(t, err)
+		require.True(t, closer.closed.Load(), "reader should be closed after synchronous extraction")
+	})
+}
+
+func TestDocumentMaxFileSize(t *testing.T) {
+	logger := mlog.CreateConsoleTestLogger(t)
+
+	data, err := testutils.ReadTestFile("sample-doc.docx")
+	require.NoError(t, err)
+
+	t.Run("a generous limit extracts the document content", func(t *testing.T) {
+		text, err := Extract(logger, "sample-doc.docx", bytes.NewReader(data), ExtractSettings{MaxFileSize: 10 * 1024 * 1024})
+		require.NoError(t, err)
+		assert.Contains(t, text, "simple")
+	})
+
+	t.Run("a tiny limit prevents the document content from being extracted", func(t *testing.T) {
+		text, err := Extract(logger, "sample-doc.docx", bytes.NewReader(data), ExtractSettings{MaxFileSize: 16})
+		require.NoError(t, err)
+		assert.NotContains(t, text, "simple")
+	})
+}
+
 func TestArchiveMaxFileSize(t *testing.T) {
 	t.Parallel()
 

server/platform/services/docextractor/documents.go

@@ -10,6 +10,8 @@ import (
 	"strings"
 
 	"code.sajari.com/docconv/v2"
+
+	"github.com/mattermost/mattermost/server/v8/channels/utils"
 )
 
 type documentExtractor struct{}
@@ -36,7 +38,7 @@ func (de *documentExtractor) Match(filename string) bool {
 	return ok
 }
 
-func (de *documentExtractor) Extract(filename string, r io.ReadSeeker, _ int64) (out string, outErr error) {
+func (de *documentExtractor) Extract(filename string, r io.ReadSeeker, maxFileSize int64) (out string, outErr error) {
 	defer func() {
 		if r := recover(); r != nil {
 			out = ""
@@ -50,7 +52,14 @@ func (de *documentExtractor) Extract(filename string, r io.ReadSeeker, _ int64)
 		return "", errors.New("unknown converter")
 	}
 
-	text, _, err := converter(r)
+	// Bound how much data the converter is allowed to read so a small upload
+	// cannot expand into an unbounded amount of in-memory work.
+	var reader io.Reader = r
+	if maxFileSize > 0 {
+		reader = utils.NewLimitedReaderWithError(r, maxFileSize)
+	}
+
+	text, _, err := converter(reader)
 	if err != nil {
 		return "", err
 	}

server/platform/services/docextractor/pdf.go

@@ -13,6 +13,8 @@ import (
 	"strings"
 
 	"github.com/ledongthuc/pdf"
+
+	"github.com/mattermost/mattermost/server/v8/channels/utils"
 )
 
 type pdfExtractor struct{}
@@ -29,7 +31,7 @@ func (pe *pdfExtractor) Match(filename string) bool {
 	return supportedExtensions[extension]
 }
 
-func (pe *pdfExtractor) Extract(filename string, r io.ReadSeeker, _ int64) (out string, outErr error) {
+func (pe *pdfExtractor) Extract(filename string, r io.ReadSeeker, maxFileSize int64) (out string, outErr error) {
 	defer func() {
 		if r := recover(); r != nil {
 			out = ""
@@ -42,7 +44,14 @@ func (pe *pdfExtractor) Extract(filename string, r io.ReadSeeker, _ int64) (out
 	}
 	defer f.Close()
 	defer os.Remove(f.Name())
-	size, err := io.Copy(f, r)
+
+	// Bound how much data is copied to disk so a small upload cannot expand
+	// into an unbounded amount of temporary storage.
+	var src io.Reader = r
+	if maxFileSize > 0 {
+		src = utils.NewLimitedReaderWithError(r, maxFileSize)
+	}
+	size, err := io.Copy(f, src)
 	if err != nil {
 		return "", fmt.Errorf("error copying data into temporary file: %v", err)
 	}

server/platform/services/docextractor/pdf_test.go

@@ -50,3 +50,30 @@ func TestWrongPdfFile(t *testing.T) {
 	_, err = extractor.Extract("sample-doc.pdf", bytes.NewReader(content), 0)
 	require.Error(t, err)
 }
+
+func TestPdfMaxFileSize(t *testing.T) {
+	extractor := pdfExtractor{}
+	content, err := testutils.ReadTestFile("sample-doc.pdf")
+	require.NoError(t, err)
+	require.Greater(t, len(content), 16, "fixture must be larger than the tight limit under test")
+
+	t.Run("a zero limit means unlimited and extracts the content", func(t *testing.T) {
+		text, err := extractor.Extract("sample-doc.pdf", bytes.NewReader(content), 0)
+		require.NoError(t, err)
+		require.Contains(t, text, "simple")
+	})
+
+	t.Run("a generous limit extracts the content", func(t *testing.T) {
+		text, err := extractor.Extract("sample-doc.pdf", bytes.NewReader(content), 10*1024*1024)
+		require.NoError(t, err)
+		require.Contains(t, text, "simple")
+	})
+
+	t.Run("a tight limit prevents extraction", func(t *testing.T) {
+		// The reader errors once it reads past the limit, so io.Copy to the
+		// temp file fails and no text is extracted.
+		text, err := extractor.Extract("sample-doc.pdf", bytes.NewReader(content), 16)
+		require.Error(t, err)
+		require.Empty(t, text)
+	})
+}