Update ElasticSearchSearchEngine to be stream-based

Signed-off-by: Eric Peterson <ericpeterson@spotify.com>
This commit is contained in:
Eric Peterson
2022-02-26 19:03:35 +01:00
parent 0d995e0ddd
commit 2c171166a0
8 changed files with 527 additions and 87 deletions
@@ -5,6 +5,8 @@
```ts
/// <reference types="node" />
import { BatchSearchEngineIndexer } from '@backstage/plugin-search-backend-node';
import { Client } from '@elastic/elasticsearch';
import { Config } from '@backstage/config';
import type { ConnectionOptions } from 'tls';
import { IndexableDocument } from '@backstage/search-common';
@@ -114,7 +116,7 @@ export class ElasticSearchSearchEngine implements SearchEngine {
indexPrefix,
}: ElasticSearchOptions): Promise<ElasticSearchSearchEngine>;
// (undocumented)
index(type: string, documents: IndexableDocument[]): Promise<void>;
getIndexer(type: string): Promise<ElasticSearchSearchEngineIndexer>;
newClient<T>(create: (options: ElasticSearchClientOptions) => T): T;
// (undocumented)
query(query: SearchQuery): Promise<SearchResultSet>;
@@ -127,4 +129,31 @@ export class ElasticSearchSearchEngine implements SearchEngine {
// (undocumented)
protected translator(query: SearchQuery): ConcreteElasticSearchQuery;
}
// Warning: (ae-missing-release-tag) "ElasticSearchSearchEngineIndexer" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
//
// @public (undocumented)
export class ElasticSearchSearchEngineIndexer extends BatchSearchEngineIndexer {
constructor(options: ElasticSearchSearchEngineIndexerOptions);
// (undocumented)
finalize(): Promise<void>;
// (undocumented)
index(documents: IndexableDocument[]): Promise<void>;
// (undocumented)
readonly indexName: string;
// (undocumented)
initialize(): Promise<void>;
}
// Warning: (ae-missing-release-tag) "ElasticSearchSearchEngineIndexerOptions" is exported by the package, but it is missing a release tag (@alpha, @beta, @public, or @internal)
//
// @public (undocumented)
export type ElasticSearchSearchEngineIndexerOptions = {
type: string;
indexPrefix: string;
indexSeparator: string;
alias: string;
logger: Logger_2;
elasticSearchClient: Client;
};
```
@@ -25,6 +25,7 @@
"dependencies": {
"@backstage/config": "^0.1.15",
"@backstage/search-common": "^0.2.4",
"@backstage/plugin-search-backend-node": "^0.4.7",
"@elastic/elasticsearch": "7.13.0",
"@acuris/aws-es-connection": "^2.2.0",
"aws-sdk": "^2.948.0",
@@ -15,7 +15,8 @@
*/
import { getVoidLogger } from '@backstage/backend-common';
import { Client } from '@elastic/elasticsearch';
import { ConfigReader } from '@backstage/config';
import { Client, errors } from '@elastic/elasticsearch';
import Mock from '@elastic/elasticsearch-mock';
import {
ConcreteElasticSearchQuery,
@@ -23,7 +24,7 @@ import {
ElasticSearchSearchEngine,
encodePageCursor,
} from './ElasticSearchSearchEngine';
import { ConfigReader } from '@backstage/config';
import { ElasticSearchSearchEngineIndexer } from './ElasticSearchSearchEngineIndexer';
class ElasticSearchSearchEngineForTranslatorTests extends ElasticSearchSearchEngine {
getTranslator() {
@@ -37,6 +38,16 @@ const options = {
Connection: mock.getConnection(),
};
const indexerMock = {
on: jest.fn(),
indexName: 'expected-index-name',
};
jest.mock('./ElasticSearchSearchEngineIndexer', () => ({
ElasticSearchSearchEngineIndexer: jest
.fn()
.mockImplementation(() => indexerMock),
}));
describe('ElasticSearchSearchEngine', () => {
let testSearchEngine: ElasticSearchSearchEngine;
let inspectableSearchEngine: ElasticSearchSearchEngineForTranslatorTests;
@@ -542,23 +553,67 @@ describe('ElasticSearchSearchEngine', () => {
});
});
describe('index', () => {
it('should index document', async () => {
const indexSpy = jest.spyOn(testSearchEngine, 'index');
const mockDocuments = [
{
title: 'testTerm',
text: 'testText',
location: 'test/location',
},
];
describe('indexer', () => {
it('should get indexer', async () => {
const indexer = await testSearchEngine.getIndexer('test-index');
// call index func and ensure the index func was invoked.
await testSearchEngine.index('test-index', mockDocuments);
expect(indexSpy).toHaveBeenCalled();
expect(indexSpy).toHaveBeenCalledWith('test-index', [
{ title: 'testTerm', text: 'testText', location: 'test/location' },
]);
expect(indexer).toStrictEqual(indexerMock);
expect(ElasticSearchSearchEngineIndexer).toHaveBeenCalledWith(
expect.objectContaining({
alias: 'test-index__search',
type: 'test-index',
indexPrefix: '',
indexSeparator: '-index__',
elasticSearchClient: client,
}),
);
expect(indexerMock.on).toHaveBeenCalledWith(
'error',
expect.any(Function),
);
});
describe('onError', () => {
let errorHandler: Function;
const error = new Error('some error');
beforeEach(async () => {
mock.clearAll();
await testSearchEngine.getIndexer('test-index');
errorHandler = indexerMock.on.mock.calls[0][1];
});
it('should check for and delete expected index', async () => {
const existsSpy = jest.fn().mockReturnValue('truthy value');
const deleteSpy = jest.fn().mockReturnValue({});
mock.add({ method: 'HEAD', path: '/expected-index-name' }, existsSpy);
mock.add({ method: 'DELETE', path: '/expected-index-name' }, deleteSpy);
await errorHandler(error);
// Check and delete HTTP requests were made.
expect(existsSpy).toHaveBeenCalled();
expect(deleteSpy).toHaveBeenCalled();
});
it('should not delete index if none exists', async () => {
// Exists call returns 404 on no index.
const existsSpy = jest.fn().mockReturnValue(
new errors.ResponseError({
statusCode: 404,
body: { status: 404 },
} as unknown as any),
);
const deleteSpy = jest.fn().mockReturnValue({});
mock.add({ method: 'HEAD', path: '/expected-index-name' }, existsSpy);
mock.add({ method: 'DELETE', path: '/expected-index-name' }, deleteSpy);
await errorHandler(error);
// Check request was made, but no delete request was made.
expect(existsSpy).toHaveBeenCalled();
expect(deleteSpy).not.toHaveBeenCalled();
});
});
});
@@ -29,8 +29,8 @@ import { Client } from '@elastic/elasticsearch';
import esb from 'elastic-builder';
import { isEmpty, isNaN as nan, isNumber } from 'lodash';
import { Logger } from 'winston';
import type { ElasticSearchClientOptions } from './ElasticSearchClientOptions';
import { ElasticSearchSearchEngineIndexer } from './ElasticSearchSearchEngineIndexer';
export type { ElasticSearchClientOptions };
@@ -58,12 +58,6 @@ type ElasticSearchResult = {
_source: IndexableDocument;
};
function duration(startTimestamp: [number, number]): string {
const delta = process.hrtime(startTimestamp);
const seconds = delta[0] + delta[1] / 1e9;
return `${seconds.toFixed(1)}s`;
}
function isBlank(str: string) {
return (isEmpty(str) && !isNumber(str)) || nan(str);
}
@@ -165,67 +159,37 @@ export class ElasticSearchSearchEngine implements SearchEngine {
this.translator = translator;
}
async index(type: string, documents: IndexableDocument[]): Promise<void> {
this.logger.info(
`Started indexing ${documents.length} documents for index ${type}`,
);
const startTimestamp = process.hrtime();
async getIndexer(type: string) {
const alias = this.constructSearchAlias(type);
const index = this.constructIndexName(type, `${Date.now()}`);
try {
const aliases = await this.elasticSearchClient.cat.aliases({
format: 'json',
name: alias,
});
const removableIndices = aliases.body.map(
(r: Record<string, any>) => r.index,
);
const indexer = new ElasticSearchSearchEngineIndexer({
type,
indexPrefix: this.indexPrefix,
indexSeparator: this.indexSeparator,
alias,
elasticSearchClient: this.elasticSearchClient,
logger: this.logger,
});
await this.elasticSearchClient.indices.create({
index,
});
const result = await this.elasticSearchClient.helpers.bulk({
datasource: documents,
onDocument() {
return {
index: { _index: index },
};
},
refreshOnCompletion: index,
});
this.logger.info(
`Indexing completed for index ${type} in ${duration(startTimestamp)}`,
result,
);
await this.elasticSearchClient.indices.updateAliases({
body: {
actions: [
{ remove: { index: this.constructIndexName(type, '*'), alias } },
{ add: { index, alias } },
],
},
});
this.logger.info('Removing stale search indices', removableIndices);
if (removableIndices.length) {
await this.elasticSearchClient.indices.delete({
index: removableIndices,
});
}
} catch (e) {
// Attempt cleanup upon failure.
indexer.on('error', async e => {
this.logger.error(`Failed to index documents for type ${type}`, e);
const response = await this.elasticSearchClient.indices.exists({
index,
});
const indexCreated = response.body;
if (indexCreated) {
this.logger.info(`Removing created index ${index}`);
await this.elasticSearchClient.indices.delete({
index,
try {
const response = await this.elasticSearchClient.indices.exists({
index: indexer.indexName,
});
const indexCreated = response.body;
if (indexCreated) {
this.logger.info(`Removing created index ${indexer.indexName}`);
await this.elasticSearchClient.indices.delete({
index: indexer.indexName,
});
}
} catch (error) {
this.logger.error(`Unable to clean up elastic index: ${error}`);
}
}
});
return indexer;
}
async query(query: SearchQuery): Promise<SearchResultSet> {
@@ -268,10 +232,6 @@ export class ElasticSearchSearchEngine implements SearchEngine {
private readonly indexSeparator = '-index__';
private constructIndexName(type: string, postFix: string) {
return `${this.indexPrefix}${type}${this.indexSeparator}${postFix}`;
}
private getTypeFromIndex(index: string) {
return index
.substring(this.indexPrefix.length)
@@ -0,0 +1,211 @@
/*
* Copyright 2022 The Backstage Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { getVoidLogger } from '@backstage/backend-common';
import { TestPipeline } from '@backstage/plugin-search-backend-node';
import { Client } from '@elastic/elasticsearch';
import Mock from '@elastic/elasticsearch-mock';
import { range } from 'lodash';
import { ElasticSearchSearchEngineIndexer } from './ElasticSearchSearchEngineIndexer';
const mock = new Mock();
const client = new Client({
node: 'http://localhost:9200',
Connection: mock.getConnection(),
});
describe('ElasticSearchSearchEngineIndexer', () => {
let indexer: ElasticSearchSearchEngineIndexer;
let bulkSpy: jest.Mock;
let catSpy: jest.Mock;
let createSpy: jest.Mock;
let aliasesSpy: jest.Mock;
let deleteSpy: jest.Mock;
beforeEach(() => {
// Instantiate the indexer to be tested.
indexer = new ElasticSearchSearchEngineIndexer({
type: 'some-type',
indexPrefix: '',
indexSeparator: '-index__',
alias: 'some-type-index__search',
logger: getVoidLogger(),
elasticSearchClient: client,
});
// Set up all requisite Elastic mocks.
mock.clearAll();
bulkSpy = jest.fn().mockReturnValue({ took: 9, errors: false, items: [] });
mock.add(
{
method: 'POST',
path: '/_bulk',
},
bulkSpy,
);
mock.add(
{
method: 'GET',
path: '/:index/_refresh',
},
jest.fn().mockReturnValue({}),
);
catSpy = jest.fn().mockReturnValue([
{
alias: 'some-type-index__search',
index: 'some-type-index__123tobedeleted',
filter: '-',
'routing.index': '-',
'routing.search': '-',
is_write_index: '-',
},
]);
mock.add(
{
method: 'GET',
path: '/_cat/aliases/some-type-index__search',
},
catSpy,
);
createSpy = jest.fn().mockReturnValue({
acknowledged: true,
shards_acknowledged: true,
index: 'single_index',
});
mock.add(
{
method: 'PUT',
path: '/:index',
},
createSpy,
);
aliasesSpy = jest.fn().mockReturnValue({});
mock.add(
{
method: 'POST',
path: '*',
},
aliasesSpy,
);
deleteSpy = jest.fn().mockReturnValue({});
mock.add(
{
method: 'DELETE',
path: '/some-type-index__123tobedeleted',
},
deleteSpy,
);
});
it('indexes documents', async () => {
const documents = [
{
title: 'testTerm',
text: 'testText',
location: 'test/location',
},
{
title: 'Another test',
text: 'Some more text',
location: 'test/location/2',
},
];
await TestPipeline.withSubject(indexer).withDocuments(documents).execute();
// Older indices should have been queried for.
expect(catSpy).toHaveBeenCalled();
// A new index should have been created.
const createdIndex = createSpy.mock.calls[0][0].path.slice(1);
expect(createdIndex).toContain('some-type-index__');
// Bulk helper should have been called with documents.
const bulkBody = bulkSpy.mock.calls[0][0].body;
expect(bulkBody[0]).toStrictEqual({ index: { _index: createdIndex } });
expect(bulkBody[1]).toStrictEqual(documents[0]);
expect(bulkBody[2]).toStrictEqual({ index: { _index: createdIndex } });
expect(bulkBody[3]).toStrictEqual(documents[1]);
// Alias should have been rotated.
expect(aliasesSpy).toHaveBeenCalled();
const aliasActions = aliasesSpy.mock.calls[0][0].body.actions;
expect(aliasActions[0]).toStrictEqual({
remove: { index: 'some-type-index__*', alias: 'some-type-index__search' },
});
expect(aliasActions[1]).toStrictEqual({
add: { index: createdIndex, alias: 'some-type-index__search' },
});
// Old index should be cleaned up.
expect(deleteSpy).toHaveBeenCalled();
});
it('handles bulk and batching during indexing', async () => {
const documents = range(550).map(i => ({
title: `Hello World ${i}`,
location: `location-${i}`,
// Generate large document sizes to trigger ES bulk flushing.
text: range(2000).join(', '),
}));
await TestPipeline.withSubject(indexer).withDocuments(documents).execute();
// Ensure multiple bulk requests were made.
expect(bulkSpy).toHaveBeenCalledTimes(2);
// Ensure the first and last documents were included in the payloads.
const docLocations: string[] = [
...bulkSpy.mock.calls[0][0].body.map((l: any) => l.location),
...bulkSpy.mock.calls[1][0].body.map((l: any) => l.location),
];
expect(docLocations).toContain('location-0');
expect(docLocations).toContain('location-549');
});
it('ignores cleanup when no existing indices exist', async () => {
const documents = [
{
title: 'testTerm',
text: 'testText',
location: 'test/location',
},
];
// Update initial alias cat to return nothing.
catSpy = jest.fn().mockReturnValue([]);
mock.clear({
method: 'GET',
path: '/_cat/aliases/some-type-index__search',
});
mock.add(
{
method: 'GET',
path: '/_cat/aliases/some-type-index__search',
},
catSpy,
);
await TestPipeline.withSubject(indexer).withDocuments(documents).execute();
// Final deletion shouldn't be called.
expect(deleteSpy).not.toHaveBeenCalled();
});
});
@@ -0,0 +1,176 @@
/*
* Copyright 2022 The Backstage Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
import { BatchSearchEngineIndexer } from '@backstage/plugin-search-backend-node';
import { IndexableDocument } from '@backstage/search-common';
import { Client } from '@elastic/elasticsearch';
import { Readable } from 'stream';
import { Logger } from 'winston';
export type ElasticSearchSearchEngineIndexerOptions = {
type: string;
indexPrefix: string;
indexSeparator: string;
alias: string;
logger: Logger;
elasticSearchClient: Client;
};
function duration(startTimestamp: [number, number]): string {
const delta = process.hrtime(startTimestamp);
const seconds = delta[0] + delta[1] / 1e9;
return `${seconds.toFixed(1)}s`;
}
export class ElasticSearchSearchEngineIndexer extends BatchSearchEngineIndexer {
private received: number = 0;
private processed: number = 0;
private removableIndices: string[] = [];
private readonly startTimestamp: [number, number];
private readonly type: string;
public readonly indexName: string;
private readonly indexPrefix: string;
private readonly indexSeparator: string;
private readonly alias: string;
private readonly logger: Logger;
private readonly sourceStream: Readable;
private readonly elasticSearchClient: Client;
private bulkResult: Promise<any>;
constructor(options: ElasticSearchSearchEngineIndexerOptions) {
super({ batchSize: 100 });
this.logger = options.logger;
this.startTimestamp = process.hrtime();
this.type = options.type;
this.indexPrefix = options.indexPrefix;
this.indexSeparator = options.indexSeparator;
this.indexName = this.constructIndexName(`${Date.now()}`);
this.alias = options.alias;
this.elasticSearchClient = options.elasticSearchClient;
// The ES client bulk helper supports stream-based indexing, but we have to
// supply the stream directly to it at instantiation-time. We can't supply
// this class itself, so instead, we create this inline stream instead.
this.sourceStream = new Readable({ objectMode: true });
this.sourceStream._read = () => {};
// eslint-disable-next-line consistent-this
const that = this;
// Keep a reference to the ES Bulk helper so that we can know when all
// documents have been successfully written to ES.
this.bulkResult = this.elasticSearchClient.helpers.bulk({
datasource: this.sourceStream,
onDocument() {
that.processed++;
return {
index: { _index: that.indexName },
};
},
refreshOnCompletion: that.indexName,
});
}
async initialize(): Promise<void> {
this.logger.info(`Started indexing documents for index ${this.type}`);
const aliases = await this.elasticSearchClient.cat.aliases({
format: 'json',
name: this.alias,
});
this.removableIndices = aliases.body.map(
(r: Record<string, any>) => r.index,
);
await this.elasticSearchClient.indices.create({
index: this.indexName,
});
}
async index(documents: IndexableDocument[]): Promise<void> {
await this.isReady();
documents.forEach(document => {
this.received++;
this.sourceStream.push(document);
});
}
async finalize(): Promise<void> {
// Wait for all documents to be processed.
await this.isReady();
// Close off the underlying stream connected to ES, indicating that no more
// documents will be written.
this.sourceStream.push(null);
// Wait for the bulk helper to finish processing.
const result = await this.bulkResult;
// Rotate aliases upon completion. Allow errors to bubble up so that we can
// clean up the create index.
this.logger.info(
`Indexing completed for index ${this.type} in ${duration(
this.startTimestamp,
)}`,
result,
);
await this.elasticSearchClient.indices.updateAliases({
body: {
actions: [
{
remove: { index: this.constructIndexName('*'), alias: this.alias },
},
{ add: { index: this.indexName, alias: this.alias } },
],
},
});
// If any indices are removable, remove them. Do not bubble up this error,
// as doing so would delete the now aliased index. Log instead.
if (this.removableIndices.length) {
this.logger.info('Removing stale search indices', this.removableIndices);
try {
await this.elasticSearchClient.indices.delete({
index: this.removableIndices,
});
} catch (e) {
this.logger.warn(`Failed to remove stale search indices: ${e}`);
}
}
}
/**
* Ensures that the number of documents sent over the wire to ES matches the
* number of documents this stream has received so far. This helps manage
* backpressure in other parts of the indexing pipeline.
*/
private isReady(): Promise<void> {
return new Promise(resolve => {
const interval = setInterval(() => {
if (this.received === this.processed) {
clearInterval(interval);
resolve();
}
}, 50);
});
}
private constructIndexName(postFix: string) {
return `${this.indexPrefix}${this.type}${this.indexSeparator}${postFix}`;
}
}
@@ -19,3 +19,7 @@ export type {
ConcreteElasticSearchQuery,
ElasticSearchClientOptions,
} from './ElasticSearchSearchEngine';
export type {
ElasticSearchSearchEngineIndexer,
ElasticSearchSearchEngineIndexerOptions,
} from './ElasticSearchSearchEngineIndexer';
@@ -21,4 +21,8 @@
*/
export { ElasticSearchSearchEngine } from './engines';
export type { ElasticSearchClientOptions } from './engines';
export type {
ElasticSearchClientOptions,
ElasticSearchSearchEngineIndexer,
ElasticSearchSearchEngineIndexerOptions,
} from './engines';