RIntermediateAggregationResultUtil.java

/**
 * diqube: Distributed Query Base.
 *
 * Copyright (C) 2015 Bastian Gloeckle
 *
 * This file is part of diqube.
 *
 * diqube is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */
package org.diqube.remote.cluster;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.ObjectInputStream;
import java.io.ObjectOutputStream;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.diqube.data.column.ColumnType;
import org.diqube.function.IntermediaryResult;
import org.diqube.function.aggregate.result.IntermediaryResultValueIterator;
import org.diqube.function.aggregate.result.serialization.IntermediateResultSerialization;
import org.diqube.function.aggregate.result.serialization.IntermediateResultSerializationResolver;
import org.diqube.remote.cluster.thrift.RColumnType;
import org.diqube.remote.cluster.thrift.RIntermediateAggregationResult;
import org.diqube.remote.cluster.thrift.RIntermediateAggregationResultValue;
import org.diqube.thrift.base.thrift.RValue;
import org.diqube.thrift.base.util.RValueUtil;
import org.diqube.util.SafeObjectInputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.google.common.collect.ImmutableSet;
import com.google.common.reflect.ClassPath;
import com.google.common.reflect.ClassPath.ClassInfo;

/**
 * Util for {@link RIntermediateAggregationResult}.
 * 
 * Serialization/deserialization adheres to {@link IntermediateResultSerialization}.
 *
 * @author Bastian Gloeckle
 */
public class RIntermediateAggregationResultUtil {
  private static final String ROOT_PKG = "org.diqube";

  private static final Logger logger = LoggerFactory.getLogger(RIntermediateAggregationResultUtil.class);

  private volatile static Set<String> whitelistedSerializableClassNames = null;

  /**
   * Deserialize a {@link RIntermediateAggregationResult} to a {@link IntermediaryResult}.
   * 
   * @throws IllegalArgumentException
   *           if data cannot be deserialized.
   */
  public static IntermediaryResult buildIntermediateAggregationResult(RIntermediateAggregationResult input)
      throws IllegalArgumentException {
    if (whitelistedSerializableClassNames == null)
      initialize();

    ColumnType type = null;
    if (input.isSetInputColumnType()) {
      switch (input.getInputColumnType()) {
      case LONG:
        type = ColumnType.LONG;
        break;
      case DOUBLE:
        type = ColumnType.DOUBLE;
        break;
      default:
        type = ColumnType.STRING;
        break;
      }
    }

    IntermediaryResult res = new IntermediaryResult(input.getOutputColName(), type);

    for (RIntermediateAggregationResultValue val : input.getValues()) {
      if (val.isSetValue()) {
        res.pushValue(RValueUtil.createValue(val.getValue()));
      } else {
        byte[] serialized = val.getSerialized();
        try (ByteArrayInputStream bais = new ByteArrayInputStream(serialized)) {
          try (ObjectInputStream ois = new SafeObjectInputStream(bais, whitelistedSerializableClassNames)) {
            res.pushValue(ois.readObject());
          }
        } catch (IOException | ClassNotFoundException e) {
          logger.error("Could not deserialize intermediate result", e);
          throw new IllegalArgumentException("Could not deserialize intermediate result", e);
        }
      }
    }

    return res;
  }

  /**
   * Serialize a {@link IntermediaryResult}.
   * 
   * @throws IllegalArgumentException
   *           If cannot be serialized
   */
  public static RIntermediateAggregationResult buildRIntermediateAggregationResult(IntermediaryResult input)
      throws IllegalArgumentException {
    if (whitelistedSerializableClassNames == null)
      initialize();

    RIntermediateAggregationResult res = new RIntermediateAggregationResult();
    res.setOutputColName(input.getOutputColName());
    if (input.getInputColumnType() != null) {
      switch (input.getInputColumnType()) {
      case STRING:
        res.setInputColumnType(RColumnType.STRING);
        break;
      case LONG:
        res.setInputColumnType(RColumnType.LONG);
        break;
      case DOUBLE:
        res.setInputColumnType(RColumnType.DOUBLE);
        break;
      }
    }

    List<RIntermediateAggregationResultValue> values = new ArrayList<>();
    IntermediaryResultValueIterator it = input.createValueIterator();
    while (it.hasNext()) {
      Object valueObject = it.next();

      RIntermediateAggregationResultValue resValue = new RIntermediateAggregationResultValue();

      RValue rvalue = RValueUtil.createRValue(valueObject);
      if (rvalue != null) {
        resValue.setValue(rvalue);
      } else {
        if (!whitelistedSerializableClassNames.contains(valueObject.getClass().getName()))
          // only a shallow check, but better than no check at all.
          throw new IllegalArgumentException("Class " + valueObject.getClass().getName() + " is not whitelisted.");

        try (ByteArrayOutputStream baos = new ByteArrayOutputStream()) {
          try (ObjectOutputStream oos = new ObjectOutputStream(baos)) {
            oos.writeObject(valueObject);
          }

          resValue.setSerialized(baos.toByteArray());
        } catch (IOException e) {
          logger.error("Could not serialize intermediary result", e);
          throw new IllegalArgumentException("Could not serialize intermediary result", e);
        }
      }

      values.add(resValue);
    }
    res.setValues(values);

    return res;
  }

  private synchronized static void initialize() {
    if (whitelistedSerializableClassNames != null)
      return;

    ClassPath cp;
    try {
      cp = ClassPath.from(RIntermediateAggregationResultUtil.class.getClassLoader());
    } catch (IOException e) {
      throw new RuntimeException("Could not initialize classpath scanning!", e);
    }
    ImmutableSet<ClassInfo> classInfos = cp.getTopLevelClassesRecursive(ROOT_PKG);

    whitelistedSerializableClassNames = new HashSet<>();

    for (ClassInfo classInfo : classInfos) {
      Class<?> clazz = classInfo.load();
      if (clazz.getAnnotation(IntermediateResultSerialization.class) != null) {
        if (!IntermediateResultSerializationResolver.class.isAssignableFrom(clazz)) {
          logger.warn("Class {} has {} annotation, but does not implement {}. Ignoring.", clazz.getName(),
              IntermediateResultSerialization.class.getSimpleName(),
              IntermediateResultSerializationResolver.class.getName());
          continue;
        }

        try {
          IntermediateResultSerializationResolver resolver =
              (IntermediateResultSerializationResolver) clazz.newInstance();

          resolver.resolve(cls -> {
            whitelistedSerializableClassNames.add(cls.getName());
            logger.debug("Whitelisted class {} for being de-/serialized for intermediate aggregation results", cls);
          });
        } catch (InstantiationException | IllegalAccessException e) {
          logger.warn("Could not instantiate {}. Ignoring.", clazz.getName(), e);
        }
      }
    }
  }
}