OpenSearch/server/src/test/java/org/opensearch/cluster/metadata/VirtualShardRoutingHelperTests.java at d13a28c6dd8fab5960a42fa7b45be9c7b1f89955 · atris/OpenSearch · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
/*
 * SPDX-License-Identifier: Apache-2.0
 *
 * The OpenSearch Contributors require contributions made to
 * this file be licensed under the Apache-2.0 license or a
 * compatible open source license.
 */

package org.opensearch.cluster.metadata;

import org.opensearch.Version;
import org.opensearch.common.settings.Settings;
import org.opensearch.test.OpenSearchTestCase;

import java.util.HashMap;
import java.util.Map;

public class VirtualShardRoutingHelperTests extends OpenSearchTestCase {

    public void testResolvePhysicalShardIdDefaultRangeBased() {
        int numPhysicalShards = 5;
        IndexMetadata metadata = IndexMetadata.builder("test")
            .settings(
                Settings.builder()
                    .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
                    .put(IndexMetadata.SETTING_NUMBER_OF_VIRTUAL_SHARDS, 20)
            )
            .numberOfShards(numPhysicalShards)
            .numberOfReplicas(1)
            .build();

        assertEquals(0, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 0));
        assertEquals(0, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 3));
        assertEquals(1, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 4));
        assertEquals(1, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 7));
        assertEquals(4, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 19));
    }

    public void testResolvePhysicalShardIdWithOverrides() {
        int numPhysicalShards = 5;
        Map<String, String> overrides = new HashMap<>();
        overrides.put("7", "1");
        overrides.put("8", "2");

        IndexMetadata.Builder builder = IndexMetadata.builder("test")
            .settings(
                Settings.builder()
                    .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
                    .put(IndexMetadata.SETTING_NUMBER_OF_VIRTUAL_SHARDS, 20)
            )
            .numberOfShards(numPhysicalShards)
            .numberOfReplicas(1);
        builder.putCustom(VirtualShardRoutingHelper.VIRTUAL_SHARDS_CUSTOM_METADATA_KEY, overrides);

        IndexMetadata metadata = builder.build();

        assertEquals(1, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 7));
        assertEquals(2, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 8));

        assertEquals(0, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 0));
        assertEquals(2, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 9));
    }

    public void testInvalidOverridesFallBackToRangeBased() {
        int numPhysicalShards = 5;
        Map<String, String> overrides = new HashMap<>();
        overrides.put("7", "not_a_number");
        overrides.put("8", "-1");
        overrides.put("19", "5");

        IndexMetadata.Builder builder = IndexMetadata.builder("test")
            .settings(
                Settings.builder()
                    .put(IndexMetadata.SETTING_VERSION_CREATED, Version.CURRENT)
                    .put(IndexMetadata.SETTING_NUMBER_OF_VIRTUAL_SHARDS, 20)
            )
            .numberOfShards(numPhysicalShards)
            .numberOfReplicas(1);
        builder.putCustom(VirtualShardRoutingHelper.VIRTUAL_SHARDS_CUSTOM_METADATA_KEY, overrides);

        IndexMetadata metadata = builder.build();

        assertEquals(1, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 7));
        assertEquals(2, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 8));
        assertEquals(4, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 19));
    }

    public void testResolvePhysicalShardIdInvalidConfigurations() {
        int numPhysicalShards = 5;

        IndexMetadata metadataDisabled = org.mockito.Mockito.mock(IndexMetadata.class);
        org.mockito.Mockito.when(metadataDisabled.getNumberOfVirtualShards()).thenReturn(-1);
        org.mockito.Mockito.when(metadataDisabled.getNumberOfShards()).thenReturn(numPhysicalShards);

        IllegalArgumentException e1 = expectThrows(
            IllegalArgumentException.class,
            () -> VirtualShardRoutingHelper.resolvePhysicalShardId(metadataDisabled, 0)
        );
        assertTrue(e1.getMessage().contains("must be enabled and be a multiple"));

        IndexMetadata metadataInvalid = org.mockito.Mockito.mock(IndexMetadata.class);
        org.mockito.Mockito.when(metadataInvalid.getNumberOfVirtualShards()).thenReturn(13);
        org.mockito.Mockito.when(metadataInvalid.getNumberOfShards()).thenReturn(numPhysicalShards);

        IllegalArgumentException e2 = expectThrows(
            IllegalArgumentException.class,
            () -> VirtualShardRoutingHelper.resolvePhysicalShardId(metadataInvalid, 0)
        );
        assertTrue(e2.getMessage().contains("must be enabled and be a multiple"));
    }

    public void testResolvePhysicalShardIdOutOfBoundsNormalization() {
        int numPhysicalShards = 5;
        IndexMetadata metadata = org.mockito.Mockito.mock(IndexMetadata.class);
        org.mockito.Mockito.when(metadata.getNumberOfVirtualShards()).thenReturn(20);
        org.mockito.Mockito.when(metadata.getNumberOfShards()).thenReturn(numPhysicalShards);

        assertEquals(4, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, -1));

        assertEquals(0, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 20));
        assertEquals(1, VirtualShardRoutingHelper.resolvePhysicalShardId(metadata, 25));
    }
}