Close transaction on upstream error

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>

.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.

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;

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);
+  });
 });

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();
+  }
 }