Many resources are needed to download a project. Please understand that we have to compensate our server costs. Thank you in advance. Project price only 1 $
You can buy this project and download/modify it how often you want.
/*
* Copyright 2015-2023 The OpenZipkin 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.
*/
package zipkin2.elasticsearch.internal;
import com.fasterxml.jackson.core.JsonGenerator;
import com.fasterxml.jackson.core.JsonParser;
import com.fasterxml.jackson.databind.JsonNode;
import com.google.auto.value.AutoValue;
import com.linecorp.armeria.common.HttpData;
import com.linecorp.armeria.common.HttpHeaderNames;
import com.linecorp.armeria.common.HttpMethod;
import com.linecorp.armeria.common.HttpRequest;
import com.linecorp.armeria.common.HttpRequestWriter;
import com.linecorp.armeria.common.MediaType;
import com.linecorp.armeria.common.RequestContext;
import com.linecorp.armeria.common.RequestHeaders;
import com.linecorp.armeria.common.util.Exceptions;
import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufAllocator;
import io.netty.buffer.ByteBufOutputStream;
import io.netty.buffer.PooledByteBufAllocator;
import io.netty.handler.codec.http.QueryStringEncoder;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.concurrent.RejectedExecutionException;
import java.util.function.Supplier;
import zipkin2.elasticsearch.ElasticsearchStorage;
import zipkin2.elasticsearch.ElasticsearchVersion;
import zipkin2.elasticsearch.internal.client.HttpCall;
import zipkin2.elasticsearch.internal.client.HttpCall.BodyConverter;
import static zipkin2.Call.propagateIfFatal;
import static zipkin2.elasticsearch.ElasticsearchVersion.V7_0;
import static zipkin2.elasticsearch.internal.JsonSerializers.OBJECT_MAPPER;
// See https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html
// exposed to re-use for testing writes of dependency links
public final class BulkCallBuilder {
// This mapper is invoked under the assumption that bulk requests return errors even when the http
// status is success. The status codes expected to be returned were undocumented as of version 7.2
// https://www.elastic.co/guide/en/elasticsearch/reference/current/docs-bulk.html
static final BodyConverter CHECK_FOR_ERRORS = new BodyConverter() {
@Override public Void convert(JsonParser parser, Supplier contentString) {
RuntimeException toThrow = null;
try {
JsonNode root = OBJECT_MAPPER.readTree(parser);
// only throw when we know it is an error
if (!root.at("/errors").booleanValue() && !root.at("/error").isObject()) return null;
String message = root.findPath("reason").textValue();
if (message == null) message = contentString.get();
Number status = root.findPath("status").numberValue();
if (status != null && status.intValue() == 429) {
toThrow = new RejectedExecutionException(message);
} else {
toThrow = new RuntimeException(message);
}
} catch (RuntimeException | IOException possiblyParseException) { // All use of jackson throws
}
if (toThrow != null) throw toThrow;
return null;
}
@Override public String toString() {
return "CheckForErrors";
}
};
final String tag;
final boolean shouldAddType;
final HttpCall.Factory http;
final String pipeline;
final boolean waitForRefresh;
// Mutated for each call to index
final List> entries = new ArrayList<>();
public BulkCallBuilder(ElasticsearchStorage es, ElasticsearchVersion version, String tag) {
this.tag = tag;
shouldAddType = version.compareTo(V7_0) < 0;
http = Internal.instance.http(es);
pipeline = es.pipeline();
waitForRefresh = es.flushOnWrites();
}
static IndexEntry newIndexEntry(String index, String typeName, T input,
BulkIndexWriter writer) {
return new AutoValue_BulkCallBuilder_IndexEntry<>(index, typeName, input, writer);
}
@AutoValue static abstract class IndexEntry {
abstract String index();
abstract String typeName();
abstract T input();
abstract BulkIndexWriter writer();
}
public void index(String index, String typeName, T input, BulkIndexWriter writer) {
entries.add(newIndexEntry(index, typeName, input, writer));
}
/**
* Creates a bulk request when there is more than one object to store
*/
public HttpCall build() {
QueryStringEncoder urlBuilder = new QueryStringEncoder("/_bulk");
if (pipeline != null) urlBuilder.addParam("pipeline", pipeline);
if (waitForRefresh) urlBuilder.addParam("refresh", "wait_for");
ByteBufAllocator alloc = RequestContext.mapCurrent(
RequestContext::alloc, () -> PooledByteBufAllocator.DEFAULT);
HttpCall.RequestSupplier request = new BulkRequestSupplier(
entries,
shouldAddType,
RequestHeaders.of(
HttpMethod.POST, urlBuilder.toString(),
HttpHeaderNames.CONTENT_TYPE, MediaType.JSON_UTF_8),
alloc);
return http.newCall(request, CHECK_FOR_ERRORS, tag);
}
static class BulkRequestSupplier implements HttpCall.RequestSupplier {
final List> entries;
final boolean shouldAddType;
final RequestHeaders headers;
final ByteBufAllocator alloc;
BulkRequestSupplier(List> entries, boolean shouldAddType,
RequestHeaders headers, ByteBufAllocator alloc) {
this.entries = Collections.unmodifiableList(entries);
this.shouldAddType = shouldAddType;
this.headers = headers;
this.alloc = alloc;
}
@Override public RequestHeaders headers() {
return headers;
}
@Override public HttpRequest get() {
HttpRequestWriter writer = HttpRequest.streaming(headers);
writeEntry(writer, 0);
return writer;
}
// There's a high chance that the response is received before the request
// is complete. This can be a problem for BulkCallBuilder when it's sending
// streaming requests. Hence, we use backpressure, instead of buffering.
//
// Follow https://github.com/line/armeria/issues/3119 for doc updates.
private void writeEntry(HttpRequestWriter writer, int index) {
if (index == entries.size()) { // out of entries.
writer.close();
return;
}
// Write the current entry directly to the current request.
if (!writer.tryWrite(HttpData.wrap(serialize(alloc, entries.get(index), shouldAddType)))) {
// Stream aborted, no need to serialize anymore.
return;
}
// Recurse to proceed to the next entry, if any.
writer.whenConsumed().thenRun(() -> writeEntry(writer, index + 1));
}
}
static ByteBuf serialize(ByteBufAllocator alloc, IndexEntry entry,
boolean shouldAddType) {
// Fuzzily assume a general small span is 600 bytes to reduce resizing while building up the
// JSON. Any extra bytes will be released back after serializing the document.
ByteBuf document = alloc.heapBuffer(600);
ByteBuf metadata = alloc.heapBuffer(200);
try {
String id = entry.writer().writeDocument(entry.input(), new ByteBufOutputStream(document));
writeIndexMetadata(new ByteBufOutputStream(metadata), entry, id, shouldAddType);
ByteBuf payload = alloc.ioBuffer(document.readableBytes() + metadata.readableBytes() + 2);
try {
payload.writeBytes(metadata).writeByte('\n').writeBytes(document).writeByte('\n');
} catch (Throwable t) {
payload.release();
propagateIfFatal(t);
Exceptions.throwUnsafely(t);
}
return payload;
} finally {
document.release();
metadata.release();
}
}
static void writeIndexMetadata(ByteBufOutputStream sink, IndexEntry entry, String id,
boolean shouldAddType) {
try (JsonGenerator writer = JsonSerializers.jsonGenerator(sink)) {
writer.writeStartObject();
writer.writeObjectFieldStart("index");
writer.writeStringField("_index", entry.index());
// the _type parameter is needed for Elasticsearch < 6.x
if (shouldAddType) writer.writeStringField("_type", entry.typeName());
writer.writeStringField("_id", id);
writer.writeEndObject();
writer.writeEndObject();
} catch (IOException e) {
throw new AssertionError(e); // No I/O writing to a Buffer.
}
}
}
