From 87ca22ce9c970af99771a2257c3b234c861e9ba7 Mon Sep 17 00:00:00 2001
From: Eric Peterson <ericpeterson@spotify.com>
Date: Fri, 7 Apr 2023 14:59:03 +0200
Subject: [PATCH] Close transaction on upstream error

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
---
 .changeset/search-destroyer-of-worlds.md      |  7 ++++
 .../src/engines/ElasticSearchSearchEngine.ts  |  3 ++
 .../PgSearchEngineIndexer.test.ts             | 34 +++++++++++++++++++
 .../PgSearchEngine/PgSearchEngineIndexer.ts   | 28 +++++++++++++++
 4 files changed, 72 insertions(+)
 create mode 100644 .changeset/search-destroyer-of-worlds.md

diff --git a/.changeset/search-destroyer-of-worlds.md b/.changeset/search-destroyer-of-worlds.md
new file mode 100644
index 0000000000..e5d29d2ad1
--- /dev/null
+++ b/.changeset/search-destroyer-of-worlds.md
@@ -0,0 +1,7 @@
+---
+'@backstage/plugin-search-backend-module-pg': patch
+---
+
+Fixed a bug that could cause orphaned PG connections to accumulate (eventually
+exhausting available connections) when errors were encountered earlier in the
+search indexing process.
diff --git a/plugins/search-backend-module-elasticsearch/src/engines/ElasticSearchSearchEngine.ts b/plugins/search-backend-module-elasticsearch/src/engines/ElasticSearchSearchEngine.ts
index 25df3101de..3294ca9d33 100644
--- a/plugins/search-backend-module-elasticsearch/src/engines/ElasticSearchSearchEngine.ts
+++ b/plugins/search-backend-module-elasticsearch/src/engines/ElasticSearchSearchEngine.ts
@@ -285,6 +285,9 @@ export class ElasticSearchSearchEngine implements SearchEngine {
     });
 
     // Attempt cleanup upon failure.
+    // todo(@backstage/discoverability-maintainers): Consider introducing a more
+    // formal mechanism for handling such errors in BatchSearchEngineIndexer and
+    // replacing this handler with it. See: #17291
     indexer.on('error', async e => {
       indexerLogger.error(`Failed to index documents for type ${type}`, e);
       let cleanupError: Error | undefined;
diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts
index f1cb3923e1..d6369d4ac3 100644
--- a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts
+++ b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.test.ts
@@ -15,6 +15,7 @@
  */
 import { TestPipeline } from '@backstage/plugin-search-backend-node';
 import { range } from 'lodash';
+import { Transform } from 'stream';
 import { PgSearchEngineIndexer } from './PgSearchEngineIndexer';
 import { DatabaseStore } from '../database';
 
@@ -156,4 +157,37 @@ describe('PgSearchEngineIndexer', () => {
     expect(result.error).toBe(expectedError);
     expect(tx.rollback).toHaveBeenCalledWith(expectedError);
   });
+
+  it('should rollback transaction on upstream error', async () => {
+    // Given a decorator that results in an error
+    let counter = 0;
+    const expectedError = new Error('Upstream error');
+    const errorDecorator = new Transform({ objectMode: true });
+    errorDecorator._transform = (chunk, _enc, cb) => {
+      counter++;
+      if (counter > 1) {
+        cb(expectedError);
+      } else {
+        cb(undefined, chunk);
+      }
+    };
+
+    // When the decorator is run in a pipeline with the PG indexer
+    const result = await TestPipeline.fromIndexer(indexer)
+      .withDecorator(errorDecorator)
+      .withDocuments([
+        { title: 'a', text: 'a', location: '/a' },
+        { title: 'b', text: 'b', location: '/b' },
+      ])
+      .execute();
+
+    // And we allow async teardown logic to complete
+    await new Promise(resolve => setImmediate(resolve));
+
+    // Then the transaction should have been closed with the expected error.
+    expect(database.getTransaction).toHaveBeenCalledTimes(1);
+    expect(database.completeInsert).not.toHaveBeenCalled();
+    expect(result.error).toBe(expectedError);
+    expect(tx.rollback).toHaveBeenCalledWith(expectedError);
+  });
 });
diff --git a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts
index ba625bc116..d1ff771d40 100644
--- a/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts
+++ b/plugins/search-backend-module-pg/src/PgSearchEngine/PgSearchEngineIndexer.ts
@@ -92,4 +92,32 @@ export class PgSearchEngineIndexer extends BatchSearchEngineIndexer {
       throw e;
     }
   }
+
+  /**
+   * Custom handler covering the case where an error occurred somewhere else in
+   * the indexing pipeline (e.g. a collator or decorator). In such cases, the
+   * finalize method is not called, which leaves a dangling transaction and
+   * therefore an open connection to PG. This handler ensures we close the
+   * transaction and associated connection.
+   *
+   * todo(@backstage/discoverability-maintainers): Consider introducing a more
+   * formal mechanism for handling such errors in BatchSearchEngineIndexer and
+   * replacing this method with it. See: #17291
+   *
+   * @internal
+   */
+  async _destroy(error: Error | null, done: (error?: Error | null) => void) {
+    // Ignore situations where there was no error.
+    if (!error) {
+      done();
+      return;
+    }
+
+    try {
+      this.tx!.rollback(error);
+    } catch {
+      // Unlikely! It was likely rolled back earlier.
+    }
+    done();
+  }
 }