Examples with NamedWriteableAwareStreamInput org.opensearch.common.io.stream.NamedWriteableAwareStreamInput used on opensource projects

Example 1 with NamedWriteableAwareStreamInput

use of org.opensearch.common.io.stream.NamedWriteableAwareStreamInput in project OpenSearch by opensearch-project.

the class PercolatorFieldMapperTests method testImplicitlySetDefaultScriptLang.

public void testImplicitlySetDefaultScriptLang() throws Exception {
    addQueryFieldMappings();
    XContentBuilder query = jsonBuilder();
    query.startObject();
    query.startObject("script");
    if (randomBoolean()) {
        query.field("script", "return true");
    } else {
        query.startObject("script");
        query.field("source", "return true");
        query.endObject();
    }
    query.endObject();
    query.endObject();
    ParsedDocument doc = mapperService.documentMapper("doc").parse(new SourceToParse("test", "doc", "1", BytesReference.bytes(XContentFactory.jsonBuilder().startObject().rawField(fieldName, new BytesArray(Strings.toString(query)).streamInput(), query.contentType()).endObject()), XContentType.JSON));
    BytesRef querySource = doc.rootDoc().getFields(fieldType.queryBuilderField.name())[0].binaryValue();
    try (InputStream in = new ByteArrayInputStream(querySource.bytes, querySource.offset, querySource.length)) {
        try (StreamInput input = new NamedWriteableAwareStreamInput(new InputStreamStreamInput(in), writableRegistry())) {
            // Query builder's content is stored via BinaryFieldMapper, which has a custom encoding
            // to encode multiple binary values into a single binary doc values field.
            // This is the reason we need to first need to read the number of values and
            // then the length of the field value in bytes.
            input.readVInt();
            input.readVInt();
            ScriptQueryBuilder queryBuilder = (ScriptQueryBuilder) input.readNamedWriteable(QueryBuilder.class);
            assertEquals(Script.DEFAULT_SCRIPT_LANG, queryBuilder.script().getLang());
        }
    }
    query = jsonBuilder();
    query.startObject();
    query.startObject("function_score");
    query.startArray("functions");
    query.startObject();
    query.startObject("script_score");
    if (randomBoolean()) {
        query.field("script", "return true");
    } else {
        query.startObject("script");
        query.field("source", "return true");
        query.endObject();
    }
    query.endObject();
    query.endObject();
    query.endArray();
    query.endObject();
    query.endObject();
    doc = mapperService.documentMapper("doc").parse(new SourceToParse("test", "doc", "1", BytesReference.bytes(XContentFactory.jsonBuilder().startObject().rawField(fieldName, new BytesArray(Strings.toString(query)).streamInput(), query.contentType()).endObject()), XContentType.JSON));
    querySource = doc.rootDoc().getFields(fieldType.queryBuilderField.name())[0].binaryValue();
    try (InputStream in = new ByteArrayInputStream(querySource.bytes, querySource.offset, querySource.length)) {
        try (StreamInput input = new NamedWriteableAwareStreamInput(new InputStreamStreamInput(in), writableRegistry())) {
            input.readVInt();
            input.readVInt();
            FunctionScoreQueryBuilder queryBuilder = (FunctionScoreQueryBuilder) input.readNamedWriteable(QueryBuilder.class);
            ScriptScoreFunctionBuilder function = (ScriptScoreFunctionBuilder) queryBuilder.filterFunctionBuilders()[0].getScoreFunction();
            assertEquals(Script.DEFAULT_SCRIPT_LANG, function.getScript().getLang());
        }
    }
}

Example 2 with NamedWriteableAwareStreamInput

use of org.opensearch.common.io.stream.NamedWriteableAwareStreamInput in project OpenSearch by opensearch-project.

the class PercolateQueryBuilder method createStore.

static PercolateQuery.QueryStore createStore(MappedFieldType queryBuilderFieldType, QueryShardContext context) {
    Version indexVersion = context.indexVersionCreated();
    NamedWriteableRegistry registry = context.getWriteableRegistry();
    return ctx -> {
        LeafReader leafReader = ctx.reader();
        BinaryDocValues binaryDocValues = leafReader.getBinaryDocValues(queryBuilderFieldType.name());
        if (binaryDocValues == null) {
            return docId -> null;
        }
        return docId -> {
            if (binaryDocValues.advanceExact(docId)) {
                BytesRef qbSource = binaryDocValues.binaryValue();
                try (InputStream in = new ByteArrayInputStream(qbSource.bytes, qbSource.offset, qbSource.length)) {
                    try (StreamInput input = new NamedWriteableAwareStreamInput(new InputStreamStreamInput(in, qbSource.length), registry)) {
                        input.setVersion(indexVersion);
                        // Query builder's content is stored via BinaryFieldMapper, which has a custom encoding
                        // to encode multiple binary values into a single binary doc values field.
                        // This is the reason we need to first need to read the number of values and
                        // then the length of the field value in bytes.
                        int numValues = input.readVInt();
                        assert numValues == 1;
                        int valueLength = input.readVInt();
                        assert valueLength > 0;
                        QueryBuilder queryBuilder = input.readNamedWriteable(QueryBuilder.class);
                        assert in.read() == -1;
                        queryBuilder = Rewriteable.rewrite(queryBuilder, context);
                        return queryBuilder.toQuery(context);
                    }
                }
            } else {
                return null;
            }
        };
    };
}

Example 3 with NamedWriteableAwareStreamInput

use of org.opensearch.common.io.stream.NamedWriteableAwareStreamInput in project OpenSearch by opensearch-project.

the class ClusterStateDiffIT method testClusterStateDiffSerialization.

public void testClusterStateDiffSerialization() throws Exception {
    NamedWriteableRegistry namedWriteableRegistry = new NamedWriteableRegistry(ClusterModule.getNamedWriteables());
    DiscoveryNode masterNode = randomNode("master");
    DiscoveryNode otherNode = randomNode("other");
    DiscoveryNodes discoveryNodes = DiscoveryNodes.builder().add(masterNode).add(otherNode).localNodeId(masterNode.getId()).build();
    ClusterState clusterState = ClusterState.builder(new ClusterName("test")).nodes(discoveryNodes).build();
    ClusterState clusterStateFromDiffs = ClusterState.Builder.fromBytes(ClusterState.Builder.toBytes(clusterState), otherNode, namedWriteableRegistry);
    int iterationCount = randomIntBetween(10, 300);
    for (int iteration = 0; iteration < iterationCount; iteration++) {
        ClusterState previousClusterState = clusterState;
        ClusterState previousClusterStateFromDiffs = clusterStateFromDiffs;
        int changesCount = randomIntBetween(1, 4);
        ClusterState.Builder builder = null;
        for (int i = 0; i < changesCount; i++) {
            if (i > 0) {
                clusterState = builder.build();
            }
            switch(randomInt(5)) {
                case 0:
                    builder = randomNodes(clusterState);
                    break;
                case 1:
                    builder = randomRoutingTable(clusterState);
                    break;
                case 2:
                    builder = randomBlocks(clusterState);
                    break;
                case 3:
                    builder = randomClusterStateCustoms(clusterState);
                    break;
                case 4:
                    builder = randomMetadataChanges(clusterState);
                    break;
                case 5:
                    builder = randomCoordinationMetadata(clusterState);
                    break;
                default:
                    throw new IllegalArgumentException("Shouldn't be here");
            }
        }
        clusterState = builder.incrementVersion().build();
        if (randomIntBetween(0, 10) < 1) {
            // Update cluster state via full serialization from time to time
            clusterStateFromDiffs = ClusterState.Builder.fromBytes(ClusterState.Builder.toBytes(clusterState), previousClusterStateFromDiffs.nodes().getLocalNode(), namedWriteableRegistry);
        } else {
            // Update cluster states using diffs
            Diff<ClusterState> diffBeforeSerialization = clusterState.diff(previousClusterState);
            BytesStreamOutput os = new BytesStreamOutput();
            diffBeforeSerialization.writeTo(os);
            byte[] diffBytes = BytesReference.toBytes(os.bytes());
            Diff<ClusterState> diff;
            try (StreamInput input = StreamInput.wrap(diffBytes)) {
                StreamInput namedInput = new NamedWriteableAwareStreamInput(input, namedWriteableRegistry);
                diff = ClusterState.readDiffFrom(namedInput, previousClusterStateFromDiffs.nodes().getLocalNode());
                clusterStateFromDiffs = diff.apply(previousClusterStateFromDiffs);
            }
        }
        try {
            // Check non-diffable elements
            assertThat(clusterStateFromDiffs.version(), equalTo(clusterState.version()));
            assertThat(clusterStateFromDiffs.coordinationMetadata(), equalTo(clusterState.coordinationMetadata()));
            // Check nodes
            assertThat(clusterStateFromDiffs.nodes().getNodes(), equalTo(clusterState.nodes().getNodes()));
            assertThat(clusterStateFromDiffs.nodes().getLocalNodeId(), equalTo(previousClusterStateFromDiffs.nodes().getLocalNodeId()));
            assertThat(clusterStateFromDiffs.nodes().getNodes(), equalTo(clusterState.nodes().getNodes()));
            for (ObjectCursor<String> node : clusterStateFromDiffs.nodes().getNodes().keys()) {
                DiscoveryNode node1 = clusterState.nodes().get(node.value);
                DiscoveryNode node2 = clusterStateFromDiffs.nodes().get(node.value);
                assertThat(node1.getVersion(), equalTo(node2.getVersion()));
                assertThat(node1.getAddress(), equalTo(node2.getAddress()));
                assertThat(node1.getAttributes(), equalTo(node2.getAttributes()));
            }
            // Check routing table
            assertThat(clusterStateFromDiffs.routingTable().version(), equalTo(clusterState.routingTable().version()));
            assertThat(clusterStateFromDiffs.routingTable().indicesRouting(), equalTo(clusterState.routingTable().indicesRouting()));
            // Check cluster blocks
            assertThat(clusterStateFromDiffs.blocks().global(), equalTo(clusterStateFromDiffs.blocks().global()));
            assertThat(clusterStateFromDiffs.blocks().indices(), equalTo(clusterStateFromDiffs.blocks().indices()));
            assertThat(clusterStateFromDiffs.blocks().disableStatePersistence(), equalTo(clusterStateFromDiffs.blocks().disableStatePersistence()));
            // Check metadata
            assertThat(clusterStateFromDiffs.metadata().version(), equalTo(clusterState.metadata().version()));
            assertThat(clusterStateFromDiffs.metadata().clusterUUID(), equalTo(clusterState.metadata().clusterUUID()));
            assertThat(clusterStateFromDiffs.metadata().transientSettings(), equalTo(clusterState.metadata().transientSettings()));
            assertThat(clusterStateFromDiffs.metadata().persistentSettings(), equalTo(clusterState.metadata().persistentSettings()));
            assertThat(clusterStateFromDiffs.metadata().indices(), equalTo(clusterState.metadata().indices()));
            assertThat(clusterStateFromDiffs.metadata().templates(), equalTo(clusterState.metadata().templates()));
            assertThat(clusterStateFromDiffs.metadata().customs(), equalTo(clusterState.metadata().customs()));
            assertThat(clusterStateFromDiffs.metadata().equalsAliases(clusterState.metadata()), is(true));
            // JSON Serialization test - make sure that both states produce similar JSON
            assertNull(differenceBetweenMapsIgnoringArrayOrder(convertToMap(clusterStateFromDiffs), convertToMap(clusterState)));
            // Smoke test - we cannot compare bytes to bytes because some elements might get serialized in different order
            // however, serialized size should remain the same
            assertThat(ClusterState.Builder.toBytes(clusterStateFromDiffs).length, equalTo(ClusterState.Builder.toBytes(clusterState).length));
        } catch (AssertionError error) {
            logger.error("Cluster state:\n{}\nCluster state from diffs:\n{}", clusterState.toString(), clusterStateFromDiffs.toString());
            throw error;
        }
    }
    logger.info("Final cluster state:[{}]", clusterState.toString());
}

Example 4 with NamedWriteableAwareStreamInput

use of org.opensearch.common.io.stream.NamedWriteableAwareStreamInput in project OpenSearch by opensearch-project.

the class ConcurrentSeqNoVersioningIT method runLinearizabilityChecker.

@SuppressForbidden(reason = "system out is ok for a command line tool")
private static void runLinearizabilityChecker(FileInputStream fileInputStream, long primaryTerm, long seqNo) throws IOException {
    StreamInput is = new InputStreamStreamInput(Base64.getDecoder().wrap(fileInputStream));
    is = new NamedWriteableAwareStreamInput(is, createNamedWriteableRegistry());
    LinearizabilityChecker.History history = readHistory(is);
    Version initialVersion = new Version(primaryTerm, seqNo);
    boolean result = new LinearizabilityChecker().isLinearizable(new CASSequentialSpec(initialVersion), history, missingResponseGenerator());
    System.out.println(LinearizabilityChecker.visualize(new CASSequentialSpec(initialVersion), history, missingResponseGenerator()));
    System.out.println("Linearizable?: " + result);
}

Example 5 with NamedWriteableAwareStreamInput

use of org.opensearch.common.io.stream.NamedWriteableAwareStreamInput in project OpenSearch by opensearch-project.

the class AllocationCommandsTests method testSerialization.

public void testSerialization() throws Exception {
    AllocationCommands commands = new AllocationCommands(new AllocateEmptyPrimaryAllocationCommand("test", 1, "node1", true), new AllocateStalePrimaryAllocationCommand("test", 2, "node1", true), new AllocateReplicaAllocationCommand("test", 2, "node1"), new MoveAllocationCommand("test", 3, "node2", "node3"), new CancelAllocationCommand("test", 4, "node5", true));
    BytesStreamOutput bytes = new BytesStreamOutput();
    AllocationCommands.writeTo(commands, bytes);
    StreamInput in = bytes.bytes().streamInput();
    // Since the commands are named writeable we need to register them and wrap the input stream
    NamedWriteableRegistry namedWriteableRegistry = new NamedWriteableRegistry(NetworkModule.getNamedWriteables());
    in = new NamedWriteableAwareStreamInput(in, namedWriteableRegistry);
    // Now we can read them!
    AllocationCommands sCommands = AllocationCommands.readFrom(in);
    assertThat(sCommands.commands().size(), equalTo(5));
    assertThat(((AllocateEmptyPrimaryAllocationCommand) (sCommands.commands().get(0))).shardId(), equalTo(1));
    assertThat(((AllocateEmptyPrimaryAllocationCommand) (sCommands.commands().get(0))).index(), equalTo("test"));
    assertThat(((AllocateEmptyPrimaryAllocationCommand) (sCommands.commands().get(0))).node(), equalTo("node1"));
    assertThat(((AllocateEmptyPrimaryAllocationCommand) (sCommands.commands().get(0))).acceptDataLoss(), equalTo(true));
    assertThat(((AllocateStalePrimaryAllocationCommand) (sCommands.commands().get(1))).shardId(), equalTo(2));
    assertThat(((AllocateStalePrimaryAllocationCommand) (sCommands.commands().get(1))).index(), equalTo("test"));
    assertThat(((AllocateStalePrimaryAllocationCommand) (sCommands.commands().get(1))).node(), equalTo("node1"));
    assertThat(((AllocateStalePrimaryAllocationCommand) (sCommands.commands().get(1))).acceptDataLoss(), equalTo(true));
    assertThat(((AllocateReplicaAllocationCommand) (sCommands.commands().get(2))).shardId(), equalTo(2));
    assertThat(((AllocateReplicaAllocationCommand) (sCommands.commands().get(2))).index(), equalTo("test"));
    assertThat(((AllocateReplicaAllocationCommand) (sCommands.commands().get(2))).node(), equalTo("node1"));
    assertThat(((MoveAllocationCommand) (sCommands.commands().get(3))).shardId(), equalTo(3));
    assertThat(((MoveAllocationCommand) (sCommands.commands().get(3))).index(), equalTo("test"));
    assertThat(((MoveAllocationCommand) (sCommands.commands().get(3))).fromNode(), equalTo("node2"));
    assertThat(((MoveAllocationCommand) (sCommands.commands().get(3))).toNode(), equalTo("node3"));
    assertThat(((CancelAllocationCommand) (sCommands.commands().get(4))).shardId(), equalTo(4));
    assertThat(((CancelAllocationCommand) (sCommands.commands().get(4))).index(), equalTo("test"));
    assertThat(((CancelAllocationCommand) (sCommands.commands().get(4))).node(), equalTo("node5"));
    assertThat(((CancelAllocationCommand) (sCommands.commands().get(4))).allowPrimary(), equalTo(true));
}