package com.mongodb.internal.connection;
import com.mongodb.internal.session.SessionContext;
import org.bson.BsonBinaryWriter;
import org.bson.BsonBinaryWriterSettings;
import org.bson.BsonDocument;
import org.bson.BsonElement;
import org.bson.BsonWriter;
import org.bson.BsonWriterSettings;
import org.bson.FieldNameValidator;
import org.bson.codecs.BsonValueCodecProvider;
import org.bson.codecs.Codec;
import org.bson.codecs.Encoder;
import org.bson.codecs.EncoderContext;
import org.bson.codecs.configuration.CodecRegistry;
import org.bson.io.BsonOutput;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import static com.mongodb.assertions.Assertions.notNull;
import static org.bson.codecs.configuration.CodecRegistries.fromProviders;
abstract class RequestMessage {
static final AtomicInteger REQUEST_ID = new AtomicInteger(1);
static final int MESSAGE_PROLOGUE_LENGTH = 16;
private static final int DOCUMENT_HEADROOM = 16 * 1024;
private static final CodecRegistry REGISTRY = fromProviders(new BsonValueCodecProvider());
private final String collectionName;
private final MessageSettings settings;
private final int id;
private final OpCode opCode;
private EncodingMetadata encodingMetadata;
static class EncodingMetadata {
private final int firstDocumentPosition;
EncodingMetadata(final int firstDocumentPosition) {
this.firstDocumentPosition = firstDocumentPosition;
}
public int getFirstDocumentPosition() {
return firstDocumentPosition;
}
}
public static int getCurrentGlobalId() {
return REQUEST_ID.get();
}
RequestMessage(final OpCode opCode, final MessageSettings settings) {
this(null, opCode, settings);
}
RequestMessage(final OpCode opCode, final int requestId, final MessageSettings settings) {
this(null, opCode, requestId, settings);
}
RequestMessage(final String collectionName, final OpCode opCode, final MessageSettings settings) {
this(collectionName, opCode, REQUEST_ID.getAndIncrement(), settings);
}
private RequestMessage(final String collectionName, final OpCode opCode, final int requestId, final MessageSettings settings) {
this.collectionName = collectionName;
this.settings = settings;
id = requestId;
this.opCode = opCode;
}
public int getId() {
return id;
}
public OpCode getOpCode() {
return opCode;
}
public MessageSettings getSettings() {
return settings;
}
public void encode(final BsonOutput bsonOutput, final SessionContext sessionContext) {
notNull("sessionContext", sessionContext);
int messageStartPosition = bsonOutput.getPosition();
writeMessagePrologue(bsonOutput);
EncodingMetadata encodingMetadata = encodeMessageBodyWithMetadata(bsonOutput, sessionContext);
backpatchMessageLength(messageStartPosition, bsonOutput);
this.encodingMetadata = encodingMetadata;
}
public EncodingMetadata getEncodingMetadata() {
return encodingMetadata;
}
protected void writeMessagePrologue(final BsonOutput bsonOutput) {
bsonOutput.writeInt32(0);
bsonOutput.writeInt32(id);
bsonOutput.writeInt32(0);
bsonOutput.writeInt32(opCode.getValue());
}
protected abstract EncodingMetadata encodeMessageBodyWithMetadata(BsonOutput bsonOutput, SessionContext sessionContext);
protected void addDocument(final BsonDocument document, final BsonOutput bsonOutput,
final FieldNameValidator validator) {
addDocument(document, getCodec(document), EncoderContext.builder().build(), bsonOutput, validator,
settings.getMaxDocumentSize() + DOCUMENT_HEADROOM, null);
}
protected void addDocument(final BsonDocument document, final BsonOutput bsonOutput,
final FieldNameValidator validator, final List<BsonElement> extraElements) {
addDocument(document, getCodec(document), EncoderContext.builder().build(), bsonOutput, validator,
settings.getMaxDocumentSize() + DOCUMENT_HEADROOM, extraElements);
}
protected void addCollectibleDocument(final BsonDocument document, final BsonOutput bsonOutput, final FieldNameValidator validator) {
addDocument(document, getCodec(document), EncoderContext.builder().isEncodingCollectibleDocument(true).build(), bsonOutput,
validator, settings.getMaxDocumentSize(), null);
}
protected void backpatchMessageLength(final int startPosition, final BsonOutput bsonOutput) {
int messageLength = bsonOutput.getPosition() - startPosition;
bsonOutput.writeInt32(bsonOutput.getPosition() - messageLength, messageLength);
}
protected String getCollectionName() {
return collectionName;
}
@SuppressWarnings("unchecked")
Codec<BsonDocument> getCodec(final BsonDocument document) {
return (Codec<BsonDocument>) REGISTRY.get(document.getClass());
}
@SuppressWarnings("unchecked")
private <T> void addDocument(final T obj, final Encoder<T> encoder, final EncoderContext encoderContext,
final BsonOutput bsonOutput, final FieldNameValidator validator, final int maxDocumentSize,
final List<BsonElement> extraElements) {
BsonBinaryWriter bsonBinaryWriter = new BsonBinaryWriter(new BsonWriterSettings(), new BsonBinaryWriterSettings(maxDocumentSize),
bsonOutput, validator);
BsonWriter bsonWriter = extraElements == null
? bsonBinaryWriter
: new ElementExtendingBsonWriter(bsonBinaryWriter, extraElements);
encoder.encode(bsonWriter, obj, encoderContext);
}
}