com.linecorp.armeria.common.thrift.text.StructContext.java Source code

Java tutorial

Introduction

Here is the source code for com.linecorp.armeria.common.thrift.text.StructContext.java

Source

// =================================================================================================
// Copyright 2011 Twitter, Inc.
// -------------------------------------------------------------------------------------------------
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this work except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file, or 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 com.linecorp.armeria.common.thrift.text;

import java.lang.reflect.Constructor;
import java.lang.reflect.Modifier;
import java.util.HashMap;
import java.util.Map;
import java.util.Map.Entry;

import javax.annotation.Nullable;

import org.apache.thrift.TApplicationException;
import org.apache.thrift.TBase;
import org.apache.thrift.TException;
import org.apache.thrift.TFieldIdEnum;
import org.apache.thrift.meta_data.EnumMetaData;
import org.apache.thrift.meta_data.FieldMetaData;
import org.apache.thrift.meta_data.FieldValueMetaData;
import org.apache.thrift.meta_data.ListMetaData;
import org.apache.thrift.meta_data.MapMetaData;
import org.apache.thrift.meta_data.SetMetaData;
import org.apache.thrift.meta_data.StructMetaData;
import org.apache.thrift.protocol.TField;
import org.apache.thrift.protocol.TType;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.databind.JsonNode;

/**
 * A struct parsing context. Builds a map from field name to TField.
 *
 * @author Alex Roetter
 */
class StructContext extends PairContext {
    private static final Logger log = LoggerFactory.getLogger(StructContext.class);

    // When processing a given thrift struct, we need certain information
    // for every field in that struct. We store that here, in a map
    // from fieldName (a string) to a TField object describing that
    // field.
    private final Map<String, TField> fieldNameMap;

    private final Map<String, Class<?>> classMap;

    /**
     * Build the name -> TField map.
     */
    StructContext(JsonNode json) {
        this(json, getCurrentThriftMessageClass());
    }

    StructContext(JsonNode json, Class<?> clazz) {
        super(json);
        classMap = new HashMap<>();
        fieldNameMap = computeFieldNameMap(clazz);
    }

    @Override
    protected TField getTFieldByName(String name) throws TException {
        if (!fieldNameMap.containsKey(name)) {
            throw new TException("Unknown field: " + name);
        }
        return fieldNameMap.get(name);
    }

    @Override
    @Nullable
    protected Class<?> getClassByFieldName(String fieldName) {
        return classMap.get(fieldName);
    }

    /**
     * I need to know what type thrift message we are processing,
     * in order to look up fields by their field name. For example,
     * i I parse a line "count : 7", I need to know we are in a
     * StatsThriftMessage, or similar, to know that count should be
     * of type int32, and have a thrift id 1.
     *
     * <p>In order to figure this out, I assume that this method was
     * called (indirectly) by the read() method in a class T which
     * is a TBase subclass. It is called that way by thrift generated
     * code. So, I iterate backwards up the call stack, stopping
     * at the first method call which belongs to a TBase object.
     * I return the Class for that object.
     *
     * <p>One could argue this is someone fragile and error prone.
     * The alternative is to modify the thrift compiler to generate
     * code which passes class information into this (and other)
     * TProtocol objects, and that seems like a lot more work. Given
     * the low level interface of TProtocol (e.g. methods like readInt(),
     * rather than readThriftMessage()), it seems likely that a TBase
     * subclass, which has the knowledge of what fields exist, as well as
     * their types & relationships, will have to be the caller of
     * the TProtocol methods.
     *
     * <p>Note: this approach does not handle TUnion, because TUnion has its own implementation of
     * read/write and any TUnion thrift structure does not override its read and write method.
     * Thus this algorithm fail to get current specific TUnion thrift structure by reading the stack.
     * To fix this, we can track call stack of nested thrift objects on our own by overriding
     * TProtocol.writeStructBegin(), rather than relying on the stack trace.
     */
    private static Class<?> getCurrentThriftMessageClass() {
        StackTraceElement[] frames = Thread.currentThread().getStackTrace();

        for (StackTraceElement f : frames) {
            String className = f.getClassName();

            try {
                Class<?> clazz = Class.forName(className);

                // Note, we need to check
                // if the class is abstract, because abstract class does not have metaDataMap
                // if the class has no-arg constructor, because FieldMetaData.getStructMetaDataMap
                //   calls clazz.newInstance
                if (isTBase(clazz) && !isAbstract(clazz) && hasNoArgConstructor(clazz)) {
                    return clazz;
                }

                if (isTApplicationException(clazz)) {
                    return clazz;
                }
            } catch (ClassNotFoundException ex) {
                log.warn("Can't find class: " + className, ex);
            }
        }
        throw new RuntimeException("Must call (indirectly) from a TBase/TApplicationException object.");
    }

    private static boolean isTBase(Class<?> clazz) {
        return TBase.class.isAssignableFrom(clazz);
    }

    private static boolean isTApplicationException(Class<?> clazz) {
        return TApplicationException.class.isAssignableFrom(clazz);
    }

    private static boolean isAbstract(Class<?> clazz) {
        return Modifier.isAbstract(clazz.getModifiers());
    }

    private static boolean hasNoArgConstructor(Class<?> clazz) {
        Constructor<?>[] allConstructors = clazz.getConstructors();
        for (Constructor<?> ctor : allConstructors) {
            Class<?>[] pType = ctor.getParameterTypes();
            if (pType.length == 0) {
                return true;
            }
        }

        return false;
    }

    /**
     * Compute a new field name map for the current thrift message
     * we are parsing.
     */
    private Map<String, TField> computeFieldNameMap(Class<?> clazz) {
        Map<String, TField> map = new HashMap<>();

        if (isTBase(clazz)) {
            // Get the metaDataMap for this Thrift class
            @SuppressWarnings("unchecked")
            Map<? extends TFieldIdEnum, FieldMetaData> metaDataMap = FieldMetaData
                    .getStructMetaDataMap((Class<? extends TBase<?, ?>>) clazz);

            for (Entry<? extends TFieldIdEnum, FieldMetaData> e : metaDataMap.entrySet()) {
                final String fieldName = e.getKey().getFieldName();
                final FieldMetaData metaData = e.getValue();

                final FieldValueMetaData elementMetaData;
                if (metaData.valueMetaData.isContainer()) {
                    if (metaData.valueMetaData instanceof SetMetaData) {
                        elementMetaData = ((SetMetaData) metaData.valueMetaData).elemMetaData;
                    } else if (metaData.valueMetaData instanceof ListMetaData) {
                        elementMetaData = ((ListMetaData) metaData.valueMetaData).elemMetaData;
                    } else if (metaData.valueMetaData instanceof MapMetaData) {
                        elementMetaData = ((MapMetaData) metaData.valueMetaData).valueMetaData;
                    } else {
                        // Unrecognized container type, but let's still continue processing without
                        // special enum support.
                        elementMetaData = metaData.valueMetaData;
                    }
                } else {
                    elementMetaData = metaData.valueMetaData;
                }

                if (elementMetaData instanceof EnumMetaData) {
                    classMap.put(fieldName, ((EnumMetaData) elementMetaData).enumClass);
                } else if (elementMetaData instanceof StructMetaData) {
                    classMap.put(fieldName, ((StructMetaData) elementMetaData).structClass);
                }

                // Workaround a bug in the generated thrift message read()
                // method by mapping the ENUM type to the INT32 type
                // The thrift generated parsing code requires that, when expecting
                // a value of enum, we actually parse a value of type int32. The
                // generated read() method then looks up the enum value in a map.
                byte type = TType.ENUM == metaData.valueMetaData.type ? TType.I32 : metaData.valueMetaData.type;

                map.put(fieldName, new TField(fieldName, type, e.getKey().getThriftFieldId()));
            }
        } else { // TApplicationException
            map.put("message", new TField("message", (byte) 11, (short) 1));
            map.put("type", new TField("type", (byte) 8, (short) 2));
        }

        return map;
    }
}