Skip to content

Commit

Permalink
Support encoding and decoding MongoDB entity fields declared as `Obje…
Browse files Browse the repository at this point in the history
…ct` + Polish code
  • Loading branch information
JamesChenX committed Jul 11, 2024
1 parent dbe9327 commit 361733a
Show file tree
Hide file tree
Showing 3 changed files with 186 additions and 29 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,7 @@
import org.bson.codecs.EncoderContext;
import org.bson.codecs.configuration.CodecRegistry;

import im.turms.server.common.infra.lang.Pair;
import im.turms.server.common.infra.lang.PrimitiveUtil;
import im.turms.server.common.infra.lang.Quadruple;
import im.turms.server.common.infra.lang.Triple;
import im.turms.server.common.infra.logging.core.logger.Logger;
import im.turms.server.common.infra.logging.core.logger.LoggerFactory;
import im.turms.server.common.infra.serialization.DeserializationException;
Expand All @@ -52,11 +49,11 @@ public class EntityCodec<T> extends MongoCodec<T> {

private static final Logger LOGGER = LoggerFactory.getLogger(EntityCodec.class);

private static final Map<Triple<Class<?>, Class<?>, Boolean>, TurmsIterableCodec> CLASS_TO_ITERABLE_CODEC =
private static final Map<IterableCodecKey<?, ?>, TurmsIterableCodec<?, ?>> KEY_TO_ITERABLE_CODEC =
new ConcurrentHashMap<>(32);
private static final Map<Quadruple<Class<?>, Class<?>, Class<?>, Boolean>, TurmsMapCodec> CLASS_TO_MAP_CODEC =
private static final Map<MapCodecKey<?, ?>, TurmsMapCodec<?, ?>> KEY_TO_MAP_CODEC =
new ConcurrentHashMap<>(32);
private static final Map<Pair<Class<?>, Boolean>, MongoCodec<?>> CLASS_TO_ENUM_CODEC =
private static final Map<EnumCodecKey<?>, MongoCodec<?>> KEY_TO_ENUM_CODEC =
new ConcurrentHashMap<>(32);

private final CodecRegistry registry;
Expand Down Expand Up @@ -252,41 +249,63 @@ private <F> F decode(EntityField<F> field, BsonReader reader, DecoderContext dec
private <F> Codec<F> getCodec(EntityField<F> field) {
Class<?> fieldClass = field.getFieldClass();
if (Iterable.class.isAssignableFrom(fieldClass)) {
return CLASS_TO_ITERABLE_CODEC.computeIfAbsent(
Triple.of(fieldClass, field.getElementClass(), field.isEnumNumber()),
triple -> {
TurmsIterableCodec iterableCodec = new TurmsIterableCodec(
triple.first(),
triple.second(),
triple.third());
return (Codec<F>) KEY_TO_ITERABLE_CODEC.computeIfAbsent(new IterableCodecKey<>(
fieldClass,
field.getElementClass(),
field.isEnumNumber()), key -> {
TurmsIterableCodec<?, ?> iterableCodec = new TurmsIterableCodec<>(
key.iterableClass,
key.elementClass,
key.isEnumNumber);
iterableCodec.setRegistry(registry);
return iterableCodec;
});
} else if (Map.class.isAssignableFrom(fieldClass)) {
Quadruple<Class<?>, Class<?>, Class<?>, Boolean> key = Quadruple.of(fieldClass,
return (Codec<F>) KEY_TO_MAP_CODEC.computeIfAbsent(new MapCodecKey<>(
(Class) fieldClass,
field.getKeyClass(),
field.getElementClass(),
field.isEnumNumber());
return CLASS_TO_MAP_CODEC.computeIfAbsent(key, quadruple -> {
TurmsMapCodec<?, ?> mapCodec = new TurmsMapCodec(
quadruple.first(),
quadruple.second(),
quadruple.third(),
quadruple.fourth());
mapCodec.setRegistry(registry);
return mapCodec;
});
field.isEnumNumber()), key -> {
TurmsMapCodec<?, ?> mapCodec = new TurmsMapCodec(
key.ownerClass,
key.keyClass,
key.valueClass,
key.isEnumNumber);
mapCodec.setRegistry(registry);
return mapCodec;
});
} else if (fieldClass.isEnum()) {
Pair<Class<?>, Boolean> pair = Pair.of(fieldClass, field.isEnumNumber());
return (Codec<F>) CLASS_TO_ENUM_CODEC.computeIfAbsent(pair,
key -> key.second()
? new EnumNumberCodec<>((Class) key.first())
: new EnumStringCodec<>((Class) key.first()));
return (Codec<F>) KEY_TO_ENUM_CODEC.computeIfAbsent(
new EnumCodecKey<>(fieldClass, field.isEnumNumber()),
key -> key.isEnumNumber
? new EnumNumberCodec<>((Class) key.fieldClass)
: new EnumStringCodec<>((Class) key.fieldClass));
} else {
if (fieldClass.isPrimitive()) {
fieldClass = PrimitiveUtil.primitiveToWrapper(fieldClass);
}
return (Codec<F>) registry.get(fieldClass);
}
}

private record IterableCodecKey<I, E>(
Class<I> iterableClass,
Class<E> elementClass,
boolean isEnumNumber
) {
}

private record MapCodecKey<K, V>(
Class<Map<K, V>> ownerClass,
Class<K> keyClass,
Class<V> valueClass,
boolean isEnumNumber
) {
}

private record EnumCodecKey<F>(
Class<F> fieldClass,
boolean isEnumNumber
) {
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@
package im.turms.server.common.storage.mongo.codec;

import java.util.Map;
import jakarta.annotation.Nullable;

import lombok.Setter;
import org.bson.codecs.Codec;
Expand Down Expand Up @@ -51,7 +52,16 @@ public <T> Codec<T> get(Class<T> clazz, CodecRegistry registry) {
return getCodec(clazz);
}

@Nullable
public <T> MongoCodec<T> getCodec(Class<T> clazz) {
if (clazz == Object.class) {
ObjectCodec objectCodec = new ObjectCodec();
objectCodec.setRegistry(registry);
return (MongoCodec<T>) objectCodec;
} else if (clazz.getClassLoader() == null) {
// Return null if the class is loaded by the bootstrap class loader.
return null;
}
return (MongoCodec<T>) classToCodec.computeIfAbsent(clazz, key -> {
MongoCodec<T> codec = (MongoCodec<T>) createCodec(key);
codec.setRegistry(registry);
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,128 @@
/*
* Copyright (C) 2019 The Turms Project
* https://github.com/turms-im/turms
*
* 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 im.turms.server.common.storage.mongo.codec;

import java.util.ArrayList;
import java.util.Date;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.regex.Pattern;

import org.bson.BsonReader;
import org.bson.BsonType;
import org.bson.BsonWriter;
import org.bson.codecs.Codec;
import org.bson.codecs.DecoderContext;
import org.bson.codecs.EncoderContext;
import org.bson.codecs.configuration.CodecRegistry;
import org.bson.types.MaxKey;
import org.bson.types.MinKey;

/**
* @author James Chen
* @implNote Limitations: Since the codec is used for {@link Object}, we cannot know the exact type
* it should be when decoding, so we use limited types when decoding.
*/
public class ObjectCodec extends MongoCodec<Object> {

@Override
public void setRegistry(CodecRegistry registry) {
super.setRegistry(registry);
}

@Override
public Object decode(BsonReader reader, DecoderContext decoderContext) {
return switch (reader.getCurrentBsonType()) {
case DOUBLE -> reader.readDouble();
case STRING -> reader.readString();
case DOCUMENT -> {
// We use LinkedHashMap because we don't know
// if the user cares about the order or not,
// so we always use LinkedHashMap to keep the order.
Map<String, Object> map = new LinkedHashMap<>();
reader.readStartDocument();
while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
String key = reader.readName();
Object value = decode(reader, decoderContext);
map.put(key, value);
}
reader.readEndDocument();
yield map;
}
case ARRAY -> {
// For our use cases, we just use small lists in most cases,
// so 32 is a reasonable size.
List<Object> list = new ArrayList<>(32);
reader.readStartArray();
while (reader.readBsonType() != BsonType.END_OF_DOCUMENT) {
Object value = decode(reader, decoderContext);
list.add(value);
}
reader.readEndArray();
yield list;
}
case BINARY -> reader.readBinaryData()
.getData();
case UNDEFINED -> {
reader.readUndefined();
yield null;
}
case OBJECT_ID -> reader.readObjectId();
case BOOLEAN -> reader.readBoolean();
case DATE_TIME -> new Date(reader.readDateTime());
case NULL -> {
reader.readNull();
yield null;
}
case REGULAR_EXPRESSION -> Pattern.compile(reader.readRegularExpression()
.getPattern());
case DB_POINTER ->
throw new UnsupportedOperationException("DBPointer is not supported");
case JAVASCRIPT -> reader.readJavaScript();
case SYMBOL -> throw new UnsupportedOperationException("Symbol is not supported");
case JAVASCRIPT_WITH_SCOPE ->
throw new UnsupportedOperationException("JavaScriptWithScope is not supported");
case INT32 -> reader.readInt32();
case TIMESTAMP -> reader.readTimestamp();
case INT64 -> reader.readInt64();
case DECIMAL128 -> reader.readDecimal128();
case MIN_KEY -> {
reader.readMinKey();
yield new MinKey();
}
case MAX_KEY -> {
reader.readMaxKey();
yield new MaxKey();
}
case END_OF_DOCUMENT -> throw new UnsupportedOperationException(
"The end of document should be handled by the caller");
};
}

@Override
public void encode(BsonWriter writer, Object value, EncoderContext encoderContext) {
Class<?> valueClass = value.getClass();
if (Object.class == valueClass) {
throw new IllegalArgumentException("The value class must not be Object");
}
Codec valueCodec = registry.get(valueClass);
valueCodec.encode(writer, value, encoderContext);
}

}

0 comments on commit 361733a

Please sign in to comment.