JsonBuffer.java

package io.jstach.rainbowgum.json;

import io.jstach.rainbowgum.LogEncoder.Buffer;

import static io.jstach.rainbowgum.json.RawJsonWriter.COMMA;
import static io.jstach.rainbowgum.json.RawJsonWriter.SEMI;

import org.jspecify.annotations.Nullable;

import io.jstach.rainbowgum.LogEvent;
import io.jstach.rainbowgum.LogOutput;

/**
 * A buffer designed for encoding JSON efficiently.
 */
public final class JsonBuffer implements Buffer {

	/**
	 * Default initial capacity of the raw JSON byte buffer, matching what this class has
	 * always used - big enough that most events never need to grow it.
	 */
	public static final int DEFAULT_INITIAL_JSON_CAPACITY = 1024 * 8;

	/**
	 * Default initial capacity of {@link #getFormattedMessageBuilder()}. Bigger than
	 * {@link StringBuilder}'s own default of 16 since most log messages exceed that
	 * immediately, but still small relative to {@link #DEFAULT_INITIAL_JSON_CAPACITY}
	 * since this buffer is just the formatted message, not the whole encoded event.
	 */
	public static final int DEFAULT_INITIAL_MESSAGE_CAPACITY = 128;

	private final RawJsonWriter jsonWriter = new RawJsonWriter(DEFAULT_INITIAL_JSON_CAPACITY);

	private final StringBuilder formattedMessageBuilder = new StringBuilder(DEFAULT_INITIAL_MESSAGE_CAPACITY);

	private final boolean prettyPrint;

	private final ExtendedFieldPrefix extendedFieldPrefix;

	private final int maxBufferSize;

	/**
	 * A flag to indicate this field is extended which means it will be prefixed with
	 * {@link ExtendedFieldPrefix}.
	 */
	public static final int EXTENDED_F = 0x00000002;

	/**
	 * ASCII extended field prefix <code>_</code>
	 */
	static final byte DEFAULT_EXTENDED_FIELD_PREFIX = '_';

	/**
	 * Create a JSON buffer.
	 * @param prettyPrint whether or not to pretty print the JSON.
	 * @param extendedFieldPrefix prefix for extended fields.
	 * @param maxBufferSize a negative value disables {@link #isOversized()} entirely.
	 * Applied, for simplicity, as a single combined threshold across both the raw JSON
	 * byte buffer and {@link #getFormattedMessageBuilder()} rather than tracked
	 * separately per buffer.
	 */
	public JsonBuffer(boolean prettyPrint, ExtendedFieldPrefix extendedFieldPrefix, int maxBufferSize) {
		super();
		this.prettyPrint = prettyPrint;
		this.extendedFieldPrefix = extendedFieldPrefix;
		this.maxBufferSize = maxBufferSize;
	}

	@Override
	public void drain(LogOutput output, LogEvent event) {
		jsonWriter.write(output, event);
		clear();
	}

	@Override
	public void clear() {
		jsonWriter.reset();
		formattedMessageBuilder.setLength(0);
		if (isOversized()) {
			jsonWriter.shrinkTo(DEFAULT_INITIAL_JSON_CAPACITY);
			formattedMessageBuilder.trimToSize();
		}
	}

	@Override
	public boolean isOversized() {
		return maxBufferSize >= 0 && (jsonWriter.capacity() + formattedMessageBuilder.capacity()) > maxBufferSize;
	}

	/**
	 * Reusable String buffer for formatted messages.
	 * @return buffer.
	 */
	public StringBuilder getFormattedMessageBuilder() {
		return formattedMessageBuilder;
	}

	/**
	 * JSON tokens.
	 */
	public enum JSONToken {

		/**
		 * Helper for writing JSON object start: {
		 */
		OBJECT_START(RawJsonWriter.OBJECT_START),

		/**
		 * Helper for writing JSON object end: }
		 */
		OBJECT_END(RawJsonWriter.OBJECT_END),

		/**
		 * Helper for writing JSON array start: [
		 */
		ARRAY_START(RawJsonWriter.ARRAY_START),

		/**
		 * Helper for writing JSON array end: ]
		 */
		ARRAY_END(RawJsonWriter.ARRAY_END),

		/**
		 * Helper for writing comma separator: ,
		 */
		COMMA(RawJsonWriter.COMMA),

		/**
		 * Helper for writing semicolon: :
		 */
		SEMI(RawJsonWriter.SEMI),

		/**
		 * Helper for writing JSON quote: "
		 */
		QUOTE(RawJsonWriter.QUOTE),

		/**
		 * Helper for writing JSON escape: \\
		 */
		ESCAPE(RawJsonWriter.ESCAPE);

		final byte raw;

		private JSONToken(byte raw) {
			this.raw = (byte) raw;
		}

	}

	private static final byte UNDERSCORE_PREFIX = '_';

	private static final byte AT_PREFIX = '@';

	private static byte LF = '\n';

	private static byte SPACE = ' ';

	/**
	 * Extended fields are just fields that have some special prefix for things like GELF
	 * and ECS.
	 */
	public enum ExtendedFieldPrefix {

		/**
		 * "_" Underscore prefix
		 */
		UNDERSCORE(UNDERSCORE_PREFIX),
		/**
		 * "@" At symbol prefix
		 */
		AT(AT_PREFIX);

		private final byte raw;

		private ExtendedFieldPrefix(byte raw) {
			this.raw = raw;
		}

	}

	/**
	 * Writes a JSON token.
	 * @param token token not null.
	 */
	public final void write(JSONToken token) {
		jsonWriter.writeByte(token.raw);
	}

	/**
	 * Efficiently writes a line feed.
	 */
	public final void writeLineFeed() {
		jsonWriter.writeByte(LF);
	}

	/**
	 * Writes a string field.
	 * @param k field name
	 * @param v value
	 * @param index the current index for comma determination
	 * @return index + 1
	 */
	public final int write(String k, @Nullable String v, int index) {
		return write(k, v, index, 0);
	}

	/**
	 * Writes a string field.
	 * @param k field name
	 * @param v value
	 * @param index the current index for comma determination
	 * @param flag see {@link #EXTENDED_F}
	 * @return index + 1
	 */
	public final int write(String k, @Nullable String v, int index, int flag) {
		if (v == null)
			return index;
		_writeStartField(k, index, flag);
		jsonWriter.writeString(v);
		_writeEndField(flag);
		return index + 1;

	}

	/**
	 * Writes a double field.
	 * @param k field name
	 * @param v value
	 * @param index the current index for comma determination
	 * @param flag see {@link #EXTENDED_F}
	 * @return index + 1
	 */
	public final int writeDouble(String k, double v, int index, int flag) {
		_writeStartField(k, index, flag);
		jsonWriter.writeDouble(v);
		_writeEndField(flag);
		return index + 1;
	}

	/**
	 * Writes a string field.
	 * @param k field name
	 * @param v value
	 * @param index the current index for comma determination
	 * @param flag see {@link #EXTENDED_F}
	 * @return index + 1
	 */
	public final int writeInt(String k, int v, int index, int flag) {
		_writeStartField(k, index, flag);
		jsonWriter.writeInt(v);
		_writeEndField(flag);
		return index + 1;
	}

	/**
	 * Writes a long field.
	 * @param k field name
	 * @param v value
	 * @param index the current index for comma determination
	 * @param flag see {@link #EXTENDED_F}
	 * @return index + 1
	 */
	public final int writeLong(String k, long v, int index, int flag) {
		_writeStartField(k, index, flag);
		jsonWriter.writeAscii(Long.toString(v));
		_writeEndField(flag);
		return index + 1;
	}

	/**
	 * Starts a nested JSON object as a field value, e.g. <code>"key":{</code>. Fields
	 * written after this and before the matching {@link #writeObjectEnd()} should use the
	 * returned index (a fresh comma-counter local to the nested object) rather than the
	 * outer index.
	 * @param k field name
	 * @param index the current index (of the enclosing object) for comma determination
	 * @param flag see {@link #EXTENDED_F}
	 * @return 0, the starting index for fields inside the nested object
	 */
	public final int writeObjectStart(String k, int index, int flag) {
		_writeStartField(k, index, flag);
		jsonWriter.writeByte(RawJsonWriter.OBJECT_START);
		return 0;
	}

	/**
	 * Ends a nested JSON object previously started with
	 * {@link #writeObjectStart(String, int, int)}.
	 */
	public final void writeObjectEnd() {
		jsonWriter.writeByte(RawJsonWriter.OBJECT_END);
	}

	/**
	 * Starts a nested JSON array as a field value, e.g. <code>"key":[</code>. Elements
	 * written after this and before the matching {@link #writeArrayEnd()} should use
	 * {@link #writeArrayElementObjectStart(int)} (with a fresh, local index) rather than
	 * the outer index.
	 * @param k field name
	 * @param index the current index (of the enclosing object) for comma determination
	 * @param flag see {@link #EXTENDED_F}
	 * @return 0, the starting index for elements inside the nested array
	 */
	public final int writeArrayStart(String k, int index, int flag) {
		_writeStartField(k, index, flag);
		jsonWriter.writeByte(RawJsonWriter.ARRAY_START);
		return 0;
	}

	/**
	 * Ends a nested JSON array previously started with
	 * {@link #writeArrayStart(String, int, int)}.
	 */
	public final void writeArrayEnd() {
		jsonWriter.writeByte(RawJsonWriter.ARRAY_END);
	}

	/**
	 * Starts an anonymous JSON object as an array element (no field name), e.g. the
	 * <code>{</code> in <code>[{...},{...}]</code>.
	 * @param index the current index (within the array) for comma determination
	 * @return 0, the starting index for fields inside the element object
	 */
	public final int writeArrayElementObjectStart(int index) {
		_writeStartElement(index);
		jsonWriter.writeByte(RawJsonWriter.OBJECT_START);
		return 0;
	}

	/**
	 * Ends an array element object previously started with
	 * {@link #writeArrayElementObjectStart(int)}.
	 */
	public final void writeArrayElementObjectEnd() {
		jsonWriter.writeByte(RawJsonWriter.OBJECT_END);
	}

	private final void _writeStartElement(int index) {
		if (index > 0) {
			jsonWriter.writeByte(COMMA);
		}
		if (prettyPrint) {
			jsonWriter.writeByte(LF);
			jsonWriter.writeByte(SPACE);
		}
	}

	private final void _writeStartField(String k, int index, int flag) {
		_writeStartElement(index);
		if ((flag & EXTENDED_F) == EXTENDED_F) {
			/*
			 * Field names come from application controlled data (e.g. MDC keys) and must
			 * be JSON escaped just like values otherwise a key containing a quote can
			 * break out of the field name and forge sibling fields.
			 * https://github.com/jstachio/rainbowgum/issues/208
			 */
			jsonWriter.writeString(((char) extendedFieldPrefix.raw) + k);
		}
		else {
			jsonWriter.writeString(k);
		}
		jsonWriter.writeByte(SEMI);
	}

	/*
	 * Intentional no-op placeholder for a not-yet-implemented feature - flag exists so
	 * call sites don't need to change once this is implemented.
	 */
	@SuppressWarnings("UnusedVariable")
	private static final void _writeEndField(int flag) {
		// ignore for now.
	}

}