Skip to content

Commit

Permalink
Clean up proto microbenchmark setup - Add method for each case
Browse files Browse the repository at this point in the history
  • Loading branch information
finnegancarroll committed Oct 9, 2024
1 parent e7eea5a commit 5b45184
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 27 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,9 @@
import org.opensearch.common.xcontent.json.JsonXContent;
import org.opensearch.core.common.io.stream.BytesStreamInput;
import org.opensearch.core.xcontent.MediaType;
import org.opensearch.core.xcontent.MediaTypeRegistry;
import org.opensearch.core.xcontent.ToXContent;
import org.opensearch.core.xcontent.XContentBuilder;
import org.opensearch.proto.search.SearchHitsProtoDef;
import org.opensearch.proto.search.SearchHitsProtoDef.SearchHitsProto;
import org.opensearch.search.SearchHits;
import org.opensearch.transport.protobuf.SearchHitsProtobuf;

Expand All @@ -44,35 +43,84 @@
@OutputTimeUnit(TimeUnit.MILLISECONDS)
@State(Scope.Benchmark)
public class SearchHitsProtobufBenchmark {
String READ_PATH = "/tmp/testFiles";
int TEST_FILES = 400;

List<SearchHitsProtoDef.SearchHitsProto> searchHits = new ArrayList<>();
// Setup proto and native representations of SearchHits
List<SearchHitsProto> searchHitsProtoItems = new ArrayList<>();
List<SearchHits> searchHitsNativePojoItems = new ArrayList<>();

/*
NOTE: Test items are read from disk. Randomized SearchHits are generated with SearchHitsTests.createTestItem().
To generate test items and run all microbenchmarks in this class:
./gradlew server:test --tests "org.opensearch.search.SearchHitsTests.testMicroBenchmarkHackGenerateTestFiles" -Dtests.security.manager=false
./gradlew -p benchmarks run --args 'SearchHitsProtobufBenchmark'
*/

@Setup
public void setup() throws IOException {
Path dir = Paths.get("/Users/carrofin/fdev/repos/OpenSearch/benchmarks/src/main/java/org/opensearch/benchmark/search/testFiles");
Path dir = Paths.get(READ_PATH);

for(int i = 0; i < 100; i++) {
for(int i = 0; i < TEST_FILES; i++) {
Path testFile = dir.resolve("testItem_" + i);
InputStream in = Files.newInputStream(testFile);
BytesStreamInput sin = new BytesStreamInput(in.readAllBytes());
SearchHitsProtobuf testItem = new SearchHitsProtobuf(new SearchHits(sin));
searchHits.add(testItem.toProto());
try (InputStream in = Files.newInputStream(testFile)) {
BytesStreamInput sin = new BytesStreamInput(in.readAllBytes());

SearchHits sh = new SearchHits(sin);
SearchHitsProto shProto = new SearchHitsProtobuf(sh).toProto();

searchHitsNativePojoItems.add(sh);
searchHitsProtoItems.add(shProto);
}
}
}

/*
./gradlew -p benchmarks run --args 'SearchHitsProtobufBenchmark.writeToNativeBench'
*/
@Benchmark
public BytesStreamOutput writeToNativeBench() throws IOException {
BytesStreamOutput bytes = new BytesStreamOutput();
for (SearchHits sh : searchHitsNativePojoItems) {
sh.writeTo(bytes);
}
return bytes;
}

// @Benchmark
// public BytesStreamOutput writeToBench() throws IOException {
// BytesStreamOutput bytes = new BytesStreamOutput();
// for (SearchHitsProtobuf shProto : searchHits) {
// shProto.writeTo(bytes);
// }
// return bytes;
// }
/*
./gradlew -p benchmarks run --args 'SearchHitsProtobufBenchmark.writeToProtoBench'
*/
@Benchmark
public BytesStreamOutput writeToProtoBench() throws IOException {
BytesStreamOutput bytes = new BytesStreamOutput();
for (SearchHitsProto shProto : searchHitsProtoItems) {
shProto.writeTo(bytes);
}
return bytes;
}

/*
./gradlew -p benchmarks run --args 'SearchHitsProtobufBenchmark.toXContNativeProtoBench'
*/
public List<XContentBuilder> toXContNativeBench() throws IOException {
List<XContentBuilder> XContList = new ArrayList<>();
for (SearchHits sh : searchHitsNativePojoItems) {
XContentBuilder builder = JsonXContent.contentBuilder();
sh.toXContent(builder, ToXContent.EMPTY_PARAMS);
XContList.add(builder);
}

return XContList;
}

/*
./gradlew -p benchmarks run --args 'SearchHitsProtobufBenchmark.toXContProtoBench'
*/
@Benchmark
public List<XContentBuilder> toXContBench() throws IOException {
public List<XContentBuilder> toXContProtoBench() throws IOException {
List<XContentBuilder> XContList = new ArrayList<>();
for (SearchHitsProtoDef.SearchHitsProto sh : searchHits) {
for (SearchHitsProto sh : searchHitsProtoItems) {
XContentBuilder builder = JsonXContent.contentBuilder();
builder.startObject();
builder.rawField("protobuf", sh.toByteString().newInput(), MediaType.fromMediaType("application/octet-stream"));
Expand Down
22 changes: 14 additions & 8 deletions server/src/test/java/org/opensearch/search/SearchHitsTests.java
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,6 @@
import org.opensearch.common.xcontent.XContentType;
import org.opensearch.common.xcontent.json.JsonXContent;
import org.opensearch.core.common.bytes.BytesReference;
import org.opensearch.core.common.io.stream.StreamOutput;
import org.opensearch.core.common.io.stream.Writeable;
import org.opensearch.core.index.Index;
import org.opensearch.core.index.shard.ShardId;
Expand All @@ -55,19 +54,28 @@
import java.io.FileOutputStream;
import java.io.IOException;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Collections;
import java.util.function.Predicate;
import java.nio.file.Files;

public class SearchHitsTests extends AbstractSerializingTestCase<SearchHits> {

public void testFinnHack () throws IOException {
// HACK HACK HACK HACK HACK HACK HACK HACK HACK
Path tmp = createTempDir();
for(int i = 0; i < 100; i++) {
/*
Hack to generate a randomized set of SearchHits objects for the purpose of benchmarking serialization ops.
See SearchHitsProtobufBenchmark.java.
Note: Security manager needs to be disabled.
f */
public void testMicroBenchmarkHackGenerateTestFiles () throws IOException {
String WRITE_PATH = "/tmp/testFiles";
int TEST_FILES = 400;

Path tmp = Files.createDirectory(Paths.get(WRITE_PATH));
for(int i = 0; i < TEST_FILES; i++) {
Path pth = tmp.resolve("testItem_" + i);
Files.createFile(pth);

SearchHits sh = createTestItem(XContentType.JSON, true, true);

try (FileOutputStream out = new FileOutputStream(pth.toFile())) {
Expand All @@ -76,8 +84,6 @@ public void testFinnHack () throws IOException {
bytes.copyBytes().writeTo(out);
}
}

System.out.println("DEBUG BREAK");
}

public static SearchHits createTestItem(boolean withOptionalInnerHits, boolean withShardTarget) {
Expand Down

0 comments on commit 5b45184

Please sign in to comment.