diff --git a/README.md b/README.md index 180c6e6c..a0bbb109 100644 --- a/README.md +++ b/README.md @@ -82,7 +82,7 @@ Make sure to commit newly created migration files. ### Indexing in ES: - `cd oclapi2/` - `docker exec -it oclapi2-api-1 python manage.py search_index --populate -f --parallel` -- for populating all indexes -- `docker exec -it oclapi2-api-1 python manage.py search_index --rebuild -f --parallel` -- for rebuild (delete and create) all indexes. +- `docker exec -it oclapi2-api-1 python manage.py search_index --rebuild -f --parallel --use-alias` -- for rebuild all indexes: builds new timestamped indexes, then swaps the aliases to them. Keep `--use-alias` once `concepts`/`mappings` are aliases: without it the rebuild deletes the indexes behind them. You can also populate/re-index specific indexes, [read more](https://django-elasticsearch-dsl.readthedocs.io/en/latest/management.html) diff --git a/core/common/management/commands/es_split_index.py b/core/common/management/commands/es_split_index.py new file mode 100644 index 00000000..b3032082 --- /dev/null +++ b/core/common/management/commands/es_split_index.py @@ -0,0 +1,158 @@ +from datetime import datetime + +from django.core.management import BaseCommand, CommandError +from django.utils import timezone +from elasticsearch_dsl.connections import connections + +# manage.py es_split_index concepts --shards 6 --dry-run preflight checks and the plan, no changes +# manage.py es_split_index concepts --shards 6 [--yes] split, then swap the alias (asks first unless --yes) +# Hard-links segments into - (writes blocked for minutes), then atomically swaps in an alias . + +REQUEST_TIMEOUT = 120 + + +class Command(BaseCommand): + client = None + help = 'Split a single-shard ES index into N primaries behind an alias of the same name.' + + def add_arguments(self, parser): + parser.add_argument('index', help='Concrete index to split, e.g. concepts') + parser.add_argument('--shards', type=int, required=True, help='Primary shards of the new index') + parser.add_argument('--dry-run', action='store_true', help='Run the preflight checks only') + parser.add_argument('--yes', action='store_true', help="Don't ask before blocking writes") + parser.add_argument('--timeout', type=int, default=3600, help='Seconds to wait for the new index to be green') + parser.add_argument( + '--max-disk-percent', type=int, default=85, + help="Abort if the shard's node could pass this disk use while the new shards merge apart") + + def handle(self, *args, **options): + index = options['index'] + shards = options['shards'] + self.client = client = connections.get_connection().options(request_timeout=REQUEST_TIMEOUT) + + store_bytes, node = self.preflight(index, shards, options['max_disk_percent']) + target = f"{index}-{datetime.now().strftime('%Y%m%d%H%M%S%f')}" + self.stdout.write( + f"Plan: block writes on '{index}' ({store_bytes / 1024 ** 3:.1f} GB on {node}), split it into '{target}' " + f"with {shards} primaries and 0 replicas pinned to {node}, check doc counts, then point alias '{index}' at " + f"'{target}' and delete '{index}'.") + if options['dry_run']: + self.stdout.write('Dry run: nothing changed.') + return + if not options['yes'] and input(f"Block writes on '{index}' and go ahead? [y/N]: ").lower() != 'y': + raise CommandError('Aborted: nothing changed.') + + blocked_at = timezone.now().isoformat() + try: + client.indices.add_block(index=index, block='write') # unlike the setting, waits for in-flight writes + self.stdout.write(f"Blocked writes on '{index}' at {blocked_at}.") + self.split(index, target, shards, node, options['timeout']) + self.check_counts(index, target) + # No client retry: a resent swap fails once the first went through, and would look like a failed swap. + client.options(max_retries=0).indices.update_aliases(actions=[ + {'add': {'index': target, 'alias': index}}, + {'remove_index': {'index': index}}, + ]) + except BaseException as ex: # pylint: disable=broad-exception-caught + swapped = self.is_swapped(index, target) + if not swapped: + self.rollback(index, target, ex, delete_target=swapped is False) + self.stderr.write(f"The swap went through although it reported: {ex}") + + self.stdout.write(self.style.SUCCESS(f"Alias '{index}' now points at '{target}'; old index deleted.")) + self.stdout.write( + f"If the indexing worker wasn't paused, re-index what changed while writes were blocked: POST " + f"/indexes/resources/{index}/ with filter={{\"updated_at__gte\": \"{blocked_at}\"}}.\n" + f"Next: POST {index}/_forcemerge?only_expunge_deletes=true and wait for merges to finish, then " + f"PUT {index}/_settings {{\"index.routing.allocation.require._name\": null}} to let shards rebalance, " + f"then PUT {index}/_settings {{\"index.number_of_replicas\": 1}}.") + + def preflight(self, index, shards, max_disk_percent): + """Returns (primary store bytes, node) once every check passes; raises CommandError otherwise.""" + client = self.client + if shards < 2: + raise CommandError('--shards must be at least 2.') + if client.indices.exists_alias(name=index): + raise CommandError(f"'{index}' is already an alias, nothing to split.") + if not client.indices.exists(index=index): + raise CommandError(f"Index '{index}' doesn't exist.") + + index_settings = client.indices.get_settings(index=index)[index]['settings']['index'] + if int(index_settings['number_of_shards']) != 1: + raise CommandError(f"'{index}' has {index_settings['number_of_shards']} primaries; only 1 is supported.") + if index_settings.get('blocks'): + raise CommandError(f"'{index}' has blocks set ({index_settings['blocks']}); clear them first.") + + health = client.cluster.health(index=index)['status'] + if health != 'green': + raise CommandError(f"'{index}' is {health}, not green.") + + primary = next( + shard for shard in client.cat.shards(index=index, format='json', bytes='b') if shard['prirep'] == 'p') + store_bytes, node = int(primary['store']), primary['node'] + allocation = next( + row for row in client.cat.allocation(format='json', bytes='b') if row['node'] == node) + # Hard-linked source files are only freed once every new shard has merged away from them. + peak_percent = 100 * (int(allocation['disk.used']) + store_bytes) / int(allocation['disk.total']) + if peak_percent > max_disk_percent: + raise CommandError( + f"{node} could reach {peak_percent:.0f}% disk while the new shards merge apart (used + " + f"{store_bytes / 1024 ** 3:.1f} GB), above {max_disk_percent}%.") + return store_bytes, node + + def split(self, index, target, shards, node, timeout): # pylint: disable=too-many-arguments + client = self.client + client.indices.flush(index=index) + client.options(request_timeout=timeout).indices.split( + index=index, target=target, settings={ + 'index.number_of_shards': shards, + 'index.number_of_replicas': 0, + 'index.blocks.write': None, + # Each new shard holds the whole source via hard links until merged; moving it early copies all that. + 'index.routing.allocation.require._name': node, + }) + self.stdout.write(f"Split '{index}' into '{target}', waiting for green...") + health = client.options(request_timeout=timeout + 30).cluster.health( + index=target, wait_for_status='green', timeout=f'{timeout}s') + if health['timed_out'] or health['status'] != 'green': + raise CommandError(f"'{target}' isn't green after {timeout}s ({health['status']}).") + + def check_counts(self, index, target): + client = self.client + client.indices.refresh(index=[index, target]) + source_count = client.count(index=index)['count'] + target_count = client.count(index=target)['count'] + if source_count != target_count: + raise CommandError(f"Doc counts differ: '{index}' {source_count}, '{target}' {target_count}.") + self.stdout.write(f'Doc counts match: {source_count}.') + + def is_swapped(self, index, target): + """Whether the alias points at the target already; None if ES couldn't be asked.""" + try: + return bool(self.client.indices.exists_alias(name=index, index=target)) + except Exception: # pylint: disable=broad-exception-caught + return None + + def rollback(self, index, target, error, delete_target): + """Deletes the target (only if it's known not to be serving) and clears the write block, then raises.""" + cleanup_errors = [] + if delete_target: + try: + self.client.indices.delete(index=target, ignore_unavailable=True) + except Exception as ex: # pylint: disable=broad-exception-caught + cleanup_errors.append(f"deleting '{target}': {ex}") + else: + cleanup_errors.append(f"couldn't tell whether alias '{index}' points at '{target}', so kept it") + try: + self.client.indices.put_settings(index=index, settings={'index.blocks.write': None}) + except Exception as ex: # pylint: disable=broad-exception-caught + cleanup_errors.append(f"clearing the write block on '{index}': {ex}") + + if cleanup_errors: + raise CommandError( + f"Split failed: {error}. Rollback incomplete ({'; '.join(cleanup_errors)}): check _cat/aliases, " + f"and the write block on '{index}'.") from error + self.stderr.write(f"Rolled back: deleted '{target}' (if created) and cleared the write block on '{index}'.") + if isinstance(error, CommandError): + raise error + raise CommandError(f'Split failed and was rolled back: {error}') from error diff --git a/core/common/models.py b/core/common/models.py index 999cef25..bdbd6a45 100644 --- a/core/common/models.py +++ b/core/common/models.py @@ -178,6 +178,15 @@ def attempt(self, start, batch, index_func): logger.error(message) ERRBIT_LOGGER.log(BatchIndexingError(message)) + def refresh(self, doc): + """Refreshes the run's index once, if it sent any docs. A failed refresh is only logged.""" + if not self.docs: + return + try: + doc._index.refresh() # pylint: disable=protected-access + except Exception as ex: # pylint: disable=broad-except + logger.warning('%s index refresh after batch indexing failed: %s', self.document_name, ex) + def finish(self): """Returns the run's summary, or raises BatchIndexingError with it if any batch failed.""" if self.failed_batches: @@ -411,8 +420,8 @@ def batch_index_full( # pylint: disable=too-many-arguments """ Full (re)index, INDEX_BATCH_SIZE docs a batch (single_batch: all in one). Every batch is attempted; if any failed, raises BatchIndexingError with the counts once they all have (see BatchIndexRun). Returns the run's - summary. - refresh=False doesn't make ES refresh after each batch (the document's auto_refresh does by default). + summary. refresh (default: the document's auto_refresh) refreshes the index once after the last batch, not + per bulk. """ if get(settings, 'TEST_MODE', False): return None @@ -424,15 +433,12 @@ def batch_index_full( # pylint: disable=too-many-arguments if select_related: queryset = queryset.select_related(*select_related) - kwargs = {} refresh = doc.django.auto_refresh if refresh is None else refresh - if refresh: # as django_elasticsearch_dsl's Document.update - kwargs['refresh'] = refresh run = BatchIndexRun(document) def index_batch(objects): run.retry_rejected(lambda: BatchIndexRun.bulk( - doc, doc._get_actions(objects, 'index'), parallel, **kwargs)) # pylint: disable=protected-access + doc, doc._get_actions(objects, 'index'), parallel)) # pylint: disable=protected-access if single_batch: run.attempt(0, list(queryset.all()), index_batch) @@ -446,6 +452,8 @@ def index_batch(objects): run.attempt(start, batch, index_batch) start += batch_size + if refresh: + run.refresh(doc) return run.finish() @staticmethod @@ -463,15 +471,12 @@ def batch_index_partial_by_ids( # pylint: disable=too-many-arguments return None doc = document() - kwargs = {} refresh = doc.django.auto_refresh if refresh is None else refresh - if refresh: - kwargs['refresh'] = refresh run = BatchIndexRun(document) def index_batch(ids): try: - run.retry_rejected(lambda: BatchIndexRun.bulk(doc, get_actions(ids), parallel, **kwargs)) + run.retry_rejected(lambda: BatchIndexRun.bulk(doc, get_actions(ids), parallel)) except BulkIndexError as err: if on_bulk_error is None: BaseModel.full_index_missing_docs_or_raise(err, queryset, document) @@ -491,6 +496,8 @@ def index_batch(ids): run.attempt(start, batch, index_batch) start += batch_size + if refresh: + run.refresh(doc) return run.finish() @staticmethod @@ -509,7 +516,7 @@ def full_index_missing_docs_or_raise(err, queryset, document, prefetch=None, sel # Docs not yet in ES -- full index so they appear with all fields BaseModel.batch_index_full( single_batch=False, queryset=queryset.filter(id__in={e['update']['_id'] for e in missing}), - document=document, prefetch=prefetch or [], select_related=select_related or [] + document=document, prefetch=prefetch or [], select_related=select_related or [], refresh=False ) missing = [] except BatchIndexingError as ex: diff --git a/core/common/tasks.py b/core/common/tasks.py index 6c92ce6e..94f0d220 100644 --- a/core/common/tasks.py +++ b/core/common/tasks.py @@ -241,10 +241,20 @@ def __run_search_index_command(command, app_names=None): if not command: return - if app_names: - call_command('search_index', f'{command}', '-f', '--models', *app_names, '--parallel') - else: - call_command('search_index', command, '-f', '--parallel') + # Builds a new timestamped index and swaps the alias to it; without this a rebuild breaks an aliased index. + extra_args = ['--use-alias'] if command == '--rebuild' else [] + # --use-alias renames the shared registry Index objects and never renames them back; this process outlives the task. + names = {index: index._name for index in registry.get_indices()} # pylint: disable=protected-access + try: + # refresh=False: no refresh per bulk chunk; the 1s refresh interval makes the docs searchable. + if app_names: + call_command( + 'search_index', f'{command}', '-f', '--models', *app_names, '--parallel', *extra_args, refresh=False) + else: + call_command('search_index', command, '-f', '--parallel', *extra_args, refresh=False) + finally: + for index, name in names.items(): + index._name = name # pylint: disable=protected-access @app.task(base=QueueOnceCustomTask, retry_kwargs={'max_retries': 0}) diff --git a/core/common/tests.py b/core/common/tests.py index 0a6d984e..f5bfb0a3 100644 --- a/core/common/tests.py +++ b/core/common/tests.py @@ -13,7 +13,7 @@ from django.conf import settings from django.contrib.auth.models import AnonymousUser, Group from django.core.files.base import File -from django.core.management import call_command +from django.core.management import call_command, CommandError from django.http import HttpResponse from django.test import TestCase, override_settings, RequestFactory from django.test.runner import DiscoverRunner @@ -1426,15 +1426,78 @@ def get_actions_batches(doc_instance): @override_settings(TEST_MODE=False) @patch('core.common.models.parallel_bulk', return_value=[]) - def test_batch_index_full_can_skip_the_forced_refresh(self, parallel_bulk_mock): - queryset, doc_instance, document = self.get_batched_full_index_mocks([[1, 2]]) + def test_batch_index_full_refreshes_once_after_the_last_batch(self, parallel_bulk_mock): + queryset, doc_instance, document = self.get_batched_full_index_mocks([[1, 2], [3]]) doc_instance.django.auto_refresh = True BaseModel.batch_index_full(False, queryset, document, None, None) - self.assertIs(parallel_bulk_mock.call_args[1]['refresh'], True) + + self.assertEqual(parallel_bulk_mock.call_count, 2) + for bulk_call in parallel_bulk_mock.call_args_list: + self.assertNotIn('refresh', bulk_call[1]) + doc_instance._index.refresh.assert_called_once_with() # pylint: disable=protected-access + + @override_settings(TEST_MODE=False) + @patch('core.common.models.parallel_bulk', return_value=[]) + def test_batch_index_full_can_skip_the_refresh(self, parallel_bulk_mock): + queryset, doc_instance, document = self.get_batched_full_index_mocks([[1, 2]]) + doc_instance.django.auto_refresh = True BaseModel.batch_index_full(False, queryset, document, None, None, refresh=False) + + self.assertNotIn('refresh', parallel_bulk_mock.call_args[1]) + doc_instance._index.refresh.assert_not_called() # pylint: disable=protected-access + + @override_settings(TEST_MODE=False) + @patch('core.common.models.parallel_bulk', return_value=[]) + def test_batch_index_full_does_not_refresh_when_nothing_was_sent(self, _): + queryset, doc_instance, document = self.get_batched_full_index_mocks([]) + doc_instance.django.auto_refresh = True + + BaseModel.batch_index_full(False, queryset, document, None, None) + + doc_instance._index.refresh.assert_not_called() # pylint: disable=protected-access + + @override_settings(TEST_MODE=False) + def test_batch_index_full_refreshes_even_when_a_batch_failed(self): + queryset, doc_instance, document = self.get_batched_full_index_mocks([[1, 2], [3]]) + doc_instance.django.auto_refresh = True + + with patch('core.common.models.parallel_bulk', side_effect=[Exception('mapping conflict'), []]), \ + patch('core.common.models.ERRBIT_LOGGER'): + with self.assertRaises(BatchIndexingError): + BaseModel.batch_index_full(False, queryset, document, None, None) + + doc_instance._index.refresh.assert_called_once_with() # pylint: disable=protected-access + + @override_settings(TEST_MODE=False) + @patch('core.common.models.parallel_bulk', return_value=[]) + def test_batch_index_full_refresh_error_does_not_fail_the_run(self, _): + queryset, doc_instance, document = self.get_batched_full_index_mocks([[1, 2]]) + doc_instance.django.auto_refresh = True + doc_instance._index.refresh.side_effect = Exception('timeout') # pylint: disable=protected-access + + summary = BaseModel.batch_index_full(False, queryset, document, None, None) + + self.assertEqual(summary['failed_batches'], 0) + self.assertEqual(summary['docs'], 2) + + @override_settings(TEST_MODE=False) + @patch('core.common.models.parallel_bulk', return_value=[]) + def test_batch_index_partial_by_ids_refreshes_once_after_the_last_batch(self, parallel_bulk_mock): + id_queryset = MagicMock() + id_queryset.__getitem__.side_effect = lambda batch_slice: [1, 2] if batch_slice.start == 0 else [] + queryset = Mock() + queryset.order_by.return_value.values_list.return_value = id_queryset + doc_instance = Mock() + doc_instance.django.auto_refresh = True + doc_instance.django.queryset_pagination = None + document = Mock(return_value=doc_instance) + + BaseModel.batch_index_partial_by_ids(queryset, document, lambda ids: iter([])) + self.assertNotIn('refresh', parallel_bulk_mock.call_args[1]) + doc_instance._index.refresh.assert_called_once_with() # pylint: disable=protected-access @override_settings(TEST_MODE=False) def test_batch_index_full_continues_remaining_batches_after_error(self): @@ -1585,8 +1648,9 @@ def test_batch_index_full_single_batch_is_one_retried_batch(self): self.assertEqual(bulk_mock.call_args_list, [call( doc_instance._get_connection.return_value, # pylint: disable=protected-access doc_instance._get_actions.return_value, # pylint: disable=protected-access - raise_on_error=False, refresh=True + raise_on_error=False )] * 2) + doc_instance._index.refresh.assert_called_once_with() # pylint: disable=protected-access self.assertEqual(self.get_actions_batches(doc_instance), [[1, 2, 3]] * 2) sleep_mock.assert_called_once_with(10) self.assertEqual(summary, {'batches': 1, 'failed_batches': 0, 'docs': 3, 'failed_docs': 0}) @@ -1741,7 +1805,8 @@ def test_batch_index_source_version_append_falls_back_to_full_index_for_missing_ queryset.filter.assert_called_once_with(id__in={11}) batch_index_full_mock.assert_called_once_with( - single_batch=False, queryset=missing_queryset, document=document, prefetch=[], select_related=[] + single_batch=False, queryset=missing_queryset, document=document, prefetch=[], select_related=[], + refresh=False ) @override_settings(TEST_MODE=False) @@ -1779,7 +1844,8 @@ def get_batch(batch_slice): self.assertEqual(bulk_mock.call_count, 2) queryset.filter.assert_called_once_with(id__in={11}) batch_index_full_mock.assert_called_once_with( - single_batch=False, queryset=missing_queryset, document=document, prefetch=[], select_related=[] + single_batch=False, queryset=missing_queryset, document=document, prefetch=[], select_related=[], + refresh=False ) @override_settings(TEST_MODE=False) @@ -1890,7 +1956,8 @@ def test_batch_index_source_version_append_indexes_missing_docs_before_failing_o queryset.filter.assert_called_once_with(id__in={'11'}) batch_index_full_mock.assert_called_once_with( - single_batch=False, queryset=missing_queryset, document=document, prefetch=[], select_related=[] + single_batch=False, queryset=missing_queryset, document=document, prefetch=[], select_related=[], + refresh=False ) sleep_mock.assert_not_called() self.assertFalse(context.exception.rejected) @@ -2156,18 +2223,44 @@ def test_handle_m2m_changed_unknown_action_noop(self): def test_populate_indexes_with_app_names(self, call_command_mock): populate_indexes(['concepts']) call_command_mock.assert_called_once_with( - 'search_index', '--populate', '-f', '--models', 'concepts', '--parallel') + 'search_index', '--populate', '-f', '--models', 'concepts', '--parallel', refresh=False) @patch('core.common.tasks.call_command') def test_populate_indexes_without_app_names(self, call_command_mock): populate_indexes(None) - call_command_mock.assert_called_once_with('search_index', '--populate', '-f', '--parallel') + call_command_mock.assert_called_once_with('search_index', '--populate', '-f', '--parallel', refresh=False) @patch('core.common.tasks.call_command') def test_rebuild_indexes_with_app_names(self, call_command_mock): rebuild_indexes(['concepts']) call_command_mock.assert_called_once_with( - 'search_index', '--rebuild', '-f', '--models', 'concepts', '--parallel') + 'search_index', '--rebuild', '-f', '--models', 'concepts', '--parallel', '--use-alias', refresh=False) + + @patch('core.common.tasks.call_command') + def test_rebuild_indexes_without_app_names(self, call_command_mock): + rebuild_indexes(None) + call_command_mock.assert_called_once_with( + 'search_index', '--rebuild', '-f', '--parallel', '--use-alias', refresh=False) + + @patch('core.common.tasks.call_command') + def test_rebuild_indexes_restores_the_index_names(self, call_command_mock): + def rename(*_, **__): # as django_elasticsearch_dsl's _rebuild does with --use-alias + ConceptDocument._index._name = 'concepts-20261006000000000000' # pylint: disable=protected-access + MappingDocument._index._name = 'mappings-20261006000000000000' # pylint: disable=protected-access + + call_command_mock.side_effect = rename + rebuild_indexes(['concepts', 'mappings']) + self.assertEqual(ConceptDocument._index._name, 'concepts') # pylint: disable=protected-access + self.assertEqual(MappingDocument._index._name, 'mappings') # pylint: disable=protected-access + + def rename_and_fail(*args, **kwargs): + rename(*args, **kwargs) + raise CommandError('populate failed') + + call_command_mock.side_effect = rename_and_fail + with self.assertRaisesMessage(CommandError, 'populate failed'): + rebuild_indexes(['concepts']) + self.assertEqual(ConceptDocument._index._name, 'concepts') # pylint: disable=protected-access @patch('core.importers.importer.Importer.run') def test_bulk_import_new(self, run_mock): @@ -3155,3 +3248,177 @@ def test_multi_token_input_expands_each_known_token(self, mock_load, mock_resolv terms = LexicalVariantDictionary.get_variant_terms('childhood leukaemia colour') self.assertEqual(set(terms), {'leukemia', 'color'}) + + +class ESSplitIndexCommandTest(OCLTestCase): + @staticmethod + def get_client(**overrides): + client = MagicMock() + client.options.return_value = client + client.indices.exists_alias.return_value = False + client.indices.exists.return_value = True + client.indices.get_settings.return_value = { + 'concepts': {'settings': {'index': {'number_of_shards': '1', 'number_of_replicas': '0'}}}} + client.cluster.health.return_value = {'status': 'green', 'timed_out': False} + client.cat.shards.return_value = [{'prirep': 'p', 'node': 'es2', 'store': str(10 * 1024 ** 3)}] + client.cat.allocation.return_value = [ + {'node': 'es', 'disk.used': '0', 'disk.total': str(10 * 1024 ** 3)}, + {'node': 'es2', 'disk.used': str(45 * 1024 ** 3), 'disk.total': str(100 * 1024 ** 3)}, + ] + client.count.return_value = {'count': 42} + for attr, value in overrides.items(): + obj = client + *path, last = attr.split('.') + for part in path: + obj = getattr(obj, part) + setattr(obj, last, value) + return client + + def run_command(self, client, *args): + stderr = Mock() + with patch('core.common.management.commands.es_split_index.connections.get_connection', return_value=client): + call_command('es_split_index', 'concepts', *args, stdout=Mock(), stderr=stderr) + return stderr + + def test_splits_then_swaps_alias_atomically(self): + client = self.get_client() + + self.run_command(client, '--shards', '6', '--yes') + + client.indices.add_block.assert_called_once_with(index='concepts', block='write') + client.indices.put_settings.assert_not_called() + client.indices.flush.assert_called_once_with(index='concepts') + split_kwargs = client.indices.split.call_args[1] + target = split_kwargs['target'] + self.assertRegex(target, r'^concepts-\d{20}$') + self.assertEqual(split_kwargs['settings'], { + 'index.number_of_shards': 6, 'index.number_of_replicas': 0, 'index.blocks.write': None, + 'index.routing.allocation.require._name': 'es2'}) + client.cluster.health.assert_called_with(index=target, wait_for_status='green', timeout='3600s') + client.options.assert_any_call(max_retries=0) + client.indices.update_aliases.assert_called_once_with(actions=[ + {'add': {'index': target, 'alias': 'concepts'}}, + {'remove_index': {'index': 'concepts'}}, + ]) + client.indices.delete.assert_not_called() + + def test_dry_run_changes_nothing(self): + client = self.get_client() + + self.run_command(client, '--shards', '6', '--dry-run') + + client.indices.add_block.assert_not_called() + client.indices.split.assert_not_called() + client.indices.update_aliases.assert_not_called() + + @patch('builtins.input', return_value='n') + def test_asks_before_blocking_writes(self, input_mock): + client = self.get_client() + + with self.assertRaisesMessage(CommandError, 'nothing changed'): + self.run_command(client, '--shards', '6') + + input_mock.assert_called_once() + client.indices.add_block.assert_not_called() + client.indices.split.assert_not_called() + + def test_preflight_rejects_unsafe_sources(self): + cases = [ + ({'indices.exists_alias.return_value': True}, 'already an alias'), + ({'indices.exists.return_value': False}, "doesn't exist"), + ({'indices.get_settings.return_value': { + 'concepts': {'settings': {'index': {'number_of_shards': '3'}}}}}, 'only 1 is supported'), + ({'indices.get_settings.return_value': {'concepts': {'settings': {'index': { + 'number_of_shards': '1', 'blocks': {'read_only_allow_delete': 'true'}}}}}}, 'blocks set'), + ({'cluster.health.return_value': {'status': 'red'}}, 'not green'), + ({'cat.allocation.return_value': [ # 76 + 10 GB of 100: past 85% while the shards merge apart + {'node': 'es2', 'disk.used': str(76 * 1024 ** 3), 'disk.total': str(100 * 1024 ** 3)}]}, + 'could reach 86% disk'), + ] + for overrides, message in cases: + client = self.get_client(**overrides) + with self.assertRaisesMessage(CommandError, message): + self.run_command(client, '--shards', '6', '--yes') + client.indices.add_block.assert_not_called() + + with self.assertRaisesMessage(CommandError, 'at least 2'): + self.run_command(self.get_client(), '--shards', '1') + + def test_rolls_back_when_counts_differ(self): + client = self.get_client() + client.count.side_effect = [{'count': 42}, {'count': 41}] + + with self.assertRaisesMessage(CommandError, 'Doc counts differ'): + self.run_command(client, '--shards', '6', '--yes') + + self.assert_rolled_back(client) + + def test_rolls_back_when_split_fails(self): + client = self.get_client() + client.indices.split.side_effect = Exception('boom') + + with self.assertRaisesMessage(CommandError, 'rolled back: boom'): + self.run_command(client, '--shards', '6', '--yes') + + self.assert_rolled_back(client) + + def test_rolls_back_when_target_never_turns_green(self): + client = self.get_client() + client.cluster.health.side_effect = [ + {'status': 'green', 'timed_out': False}, {'status': 'yellow', 'timed_out': True}] + + with self.assertRaisesMessage(CommandError, "isn't green"): + self.run_command(client, '--shards', '6', '--yes') + + self.assert_rolled_back(client) + + def test_clears_a_block_that_applied_but_timed_out(self): + client = self.get_client() + client.indices.add_block.side_effect = Exception('ConnectionTimeout') + + with self.assertRaisesMessage(CommandError, 'rolled back: ConnectionTimeout'): + self.run_command(client, '--shards', '6', '--yes') + + client.indices.split.assert_not_called() + client.indices.put_settings.assert_called_once_with(index='concepts', settings={'index.blocks.write': None}) + + def test_does_not_roll_back_a_swap_that_went_through(self): + client = self.get_client() + client.indices.update_aliases.side_effect = Exception('ConnectionTimeout') # applied, response lost + client.indices.exists_alias.side_effect = [False, True] # preflight, then the reconcile check + + stderr = self.run_command(client, '--shards', '6', '--yes') + + target = client.indices.split.call_args[1]['target'] + client.indices.exists_alias.assert_called_with(name='concepts', index=target) + client.indices.delete.assert_not_called() + client.indices.put_settings.assert_not_called() + self.assertIn('swap went through', stderr.write.call_args[0][0]) + + def test_keeps_the_target_when_the_swap_state_is_unknown(self): + client = self.get_client() + client.indices.update_aliases.side_effect = Exception('ConnectionTimeout') + client.indices.exists_alias.side_effect = [False, Exception('ES unreachable')] + + with self.assertRaisesMessage(CommandError, 'Rollback incomplete'): + self.run_command(client, '--shards', '6', '--yes') + + client.indices.delete.assert_not_called() + client.indices.put_settings.assert_called_once_with(index='concepts', settings={'index.blocks.write': None}) + + def test_clears_the_block_even_when_deleting_the_target_fails(self): + client = self.get_client() + client.indices.split.side_effect = Exception('boom') + client.indices.delete.side_effect = Exception('delete timed out') + + with self.assertRaisesMessage(CommandError, 'Rollback incomplete'): + self.run_command(client, '--shards', '6', '--yes') + + client.indices.put_settings.assert_called_once_with(index='concepts', settings={'index.blocks.write': None}) + + @staticmethod + def assert_rolled_back(client): + client.indices.update_aliases.assert_not_called() + target = client.indices.split.call_args[1]['target'] + client.indices.delete.assert_called_once_with(index=target, ignore_unavailable=True) + client.indices.put_settings.assert_called_once_with(index='concepts', settings={'index.blocks.write': None}) diff --git a/core/concepts/documents.py b/core/concepts/documents.py index d8eb0a14..10aebfce 100644 --- a/core/concepts/documents.py +++ b/core/concepts/documents.py @@ -23,7 +23,10 @@ class ConceptDocument(Document): class Index: name = 'concepts' - settings = {'number_of_shards': 1, 'number_of_replicas': 0} + settings = { + 'number_of_shards': settings.ES_CONCEPTS_SHARDS, + 'number_of_replicas': settings.ES_CONCEPTS_REPLICAS + } id = fields.TextField(attr='mnemonic') id_lowercase = fields.KeywordField(attr='mnemonic', normalizer="lowercase") diff --git a/core/concepts/embeddings.py b/core/concepts/embeddings.py index 957b9dff..c56619c9 100644 --- a/core/concepts/embeddings.py +++ b/core/concepts/embeddings.py @@ -124,9 +124,9 @@ def get_stored_vectors(self, doc_ids): @property def index_name(self): - """The concepts index, which the stored docs are read from.""" + """The concepts alias, which the stored docs are read from (not _index._name, which a rebuild renames).""" from core.concepts.documents import ConceptDocument - return ConceptDocument._index._name # pylint: disable=protected-access + return ConceptDocument.Index.name @staticmethod def as_list(value): diff --git a/core/mappings/documents.py b/core/mappings/documents.py index d600a303..6af35fcd 100644 --- a/core/mappings/documents.py +++ b/core/mappings/documents.py @@ -1,3 +1,4 @@ +from django.conf import settings from django_elasticsearch_dsl import Document, fields from django_elasticsearch_dsl.registries import registry from pydash import get @@ -10,7 +11,10 @@ class MappingDocument(Document): class Index: name = 'mappings' - settings = {'number_of_shards': 1, 'number_of_replicas': 0} + settings = { + 'number_of_shards': settings.ES_MAPPINGS_SHARDS, + 'number_of_replicas': settings.ES_MAPPINGS_REPLICAS + } class Django: model = Mapping diff --git a/core/settings.py b/core/settings.py index e9026dc6..e981dbc8 100644 --- a/core/settings.py +++ b/core/settings.py @@ -300,6 +300,11 @@ def get_set_from_env(name): # ES_BULK_RETRY_BACKOFF_SECONDS before the first retry and doubling it for each one after (defaults: 10+20+40+80s). ES_BULK_RETRY_MAX_ATTEMPTS = int(os.environ.get('ES_BULK_RETRY_MAX_ATTEMPTS', 5)) ES_BULK_RETRY_BACKOFF_SECONDS = float(os.environ.get('ES_BULK_RETRY_BACKOFF_SECONDS', 10)) +# Shard layout used only when an index is created (search_index --create / --rebuild --use-alias). +ES_CONCEPTS_SHARDS = int(os.environ.get('ES_CONCEPTS_SHARDS', 1)) +ES_CONCEPTS_REPLICAS = int(os.environ.get('ES_CONCEPTS_REPLICAS', 0)) +ES_MAPPINGS_SHARDS = int(os.environ.get('ES_MAPPINGS_SHARDS', 1)) +ES_MAPPINGS_REPLICAS = int(os.environ.get('ES_MAPPINGS_REPLICAS', 0)) # Bulk imports of up to this many lines index what they import as each part finishes, unless the request's `index` # param says otherwise. Bigger ones leave it to a release or a reindex unless asked, as indexing them saturates the # shared indexing worker for hours (and re-embeds every row of a semantic HEAD).