JsonRequest.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.ui.websocket.request;

import java.io.IOException;
import java.lang.reflect.InvocationTargetException;
import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicBoolean;

import javax.annotation.PostConstruct;
import javax.inject.Inject;
import javax.websocket.Session;

import org.apache.thrift.TServiceClient;
import org.apache.thrift.protocol.TCompactProtocol;
import org.apache.thrift.protocol.TMultiplexedProtocol;
import org.apache.thrift.protocol.TProtocol;
import org.apache.thrift.transport.TFramedTransport;
import org.apache.thrift.transport.TSocket;
import org.apache.thrift.transport.TTransport;
import org.apache.thrift.transport.TTransportException;
import org.diqube.remote.query.thrift.QueryResultService.Iface;
import org.diqube.thrift.base.thrift.AuthenticationException;
import org.diqube.thrift.base.thrift.Ticket;
import org.diqube.ui.DiqubeServletConfig;
import org.diqube.ui.UiQueryRegistry;
import org.diqube.ui.websocket.request.commands.AsyncJsonCommand;
import org.diqube.ui.websocket.request.commands.JsonCommand;
import org.diqube.ui.websocket.result.ExceptionJsonResult;
import org.diqube.ui.websocket.result.JsonResult;
import org.diqube.ui.websocket.result.JsonResultEnvelope;
import org.diqube.ui.websocket.result.JsonResultSerializer;
import org.diqube.ui.websocket.result.JsonResultSerializer.JsonPayloadSerializerException;
import org.diqube.util.Pair;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.fasterxml.jackson.annotation.JsonIgnore;

/**
 * A request that was sent by the client and we should execute - it contains a {@link JsonCommand}, and the context of
 * this request. The context is defined by the session/requestId pair which is unique per request from the client - it
 * though can happen that two commands are sent using the same session/requestId pair, in which case this class ensures
 * that the corresponding commands get the same "environment".
 * 
 * <p>
 * An example could be that first a "query" command is executed with a session/requestId pair and after that a "cancel"
 * command - the latter will then cancel the former, because they have the same environment.
 * 
 * <p>
 * The environment is basically specified by instances of {@link CommandClusterInteraction} and
 * {@link CommandResultHandler} whcih are provided by this class to the {@link JsonCommand}.
 * 
 *
 * @author Bastian Gloeckle
 */
public class JsonRequest {
  private static final Logger logger = LoggerFactory.getLogger(JsonRequest.class);

  /**
   * The requestID that was created by the client to uniquely identify this request. Note that the uniqueness is <b>not
   * global</b>, but only local to the {@link #session}. This means that different sessions might (and actually will)
   * use the same request IDs to reference different requests!
   * 
   * <p>
   * Therefore one always needs to inspect both values, {@link #session} and {@link #requestId} to globally uniquely
   * identify a request.
   */
  private String requestId;
  /** The session that can be used to send data back to the client. Sync on this object when sending! */
  private Session session;
  private JsonCommand jsonCommand;

  @Inject
  @JsonIgnore
  private DiqubeServletConfig config;

  @Inject
  @JsonIgnore
  private UiQueryRegistry queryResultRegistry;

  @Inject
  @JsonIgnore
  private JsonResultSerializer serializer;

  /** {@link Runnable}s that need to be executed to clean up. */
  @JsonIgnore
  private List<Runnable> cleanupActions = new ArrayList<>();

  /**
   * {@link CommandClusterInteraction} that is passed to the command which can safely interact with the diqube cluster
   * through this instance.
   */
  private CommandClusterInteraction commandClusterInteraction;
  private JsonRequestRegistry requestRegistry;

  private Ticket ticket;

  /**
   * @param ticket
   *          <code>null</code> or a {@link Ticket} which has been validated already.
   */
  /* package */ JsonRequest(Session session, Ticket ticket, String requestId, JsonCommand jsonCommand,
      JsonRequestRegistry requestRegistry) {
    this.session = session;
    this.ticket = ticket;
    this.requestId = requestId;
    this.jsonCommand = jsonCommand;
    this.requestRegistry = requestRegistry;
  }

  @PostConstruct
  public void initialize() {
    commandClusterInteraction = new AbstractCommandClusterInteraction(config, ticket) {
      @Override
      protected void registerQueryThriftResultCallback(Pair<String, Short> node, UUID queryUuid, Iface resultHandler) {
        queryResultRegistry.registerThriftResultCallback(session, requestId, node, queryUuid, resultHandler);

        cleanupActions.add(() -> queryResultRegistry.unregisterQuery(requestId, queryUuid));
      }

      @Override
      protected Pair<UUID, Pair<String, Short>> findQueryUuidAndServerAddr() {
        UUID queryUuid = queryResultRegistry.getQueryUuid(session, requestId);
        Pair<String, Short> node = queryResultRegistry.getDiqubeServerAddr(queryUuid);

        if (queryUuid == null || node == null)
          return null;
        return new Pair<>(queryUuid, node);
      }

      @Override
      protected <T extends TServiceClient> T openConnection(Class<? extends T> thriftClientClass, String serviceName,
          Pair<String, Short> node) {
        TTransport transport = new TFramedTransport(new TSocket(node.getLeft(), node.getRight()));
        TProtocol protocol = new TMultiplexedProtocol(new TCompactProtocol(transport), serviceName);

        T res;
        try {
          res = thriftClientClass.getConstructor(TProtocol.class).newInstance(protocol);
        } catch (InstantiationException | IllegalAccessException | IllegalArgumentException | InvocationTargetException
            | NoSuchMethodException | SecurityException e) {
          throw new RuntimeException("Could not instantiate thrift client", e);
        }

        try {
          transport.open();
        } catch (TTransportException e) {
          return null;
        }

        cleanupActions.add(() -> transport.close());

        return res;
      }
    };
  }

  /**
   * Execute the command of this request. Fully handles all interaction needed with the client.
   * 
   * It is expected that the request was registered at {@link JsonRequestRegistry}. It will unregister itself as soon as
   * it has completed.
   * 
   * If the command throws an exception, {@link #sendException(Throwable)} will be run and the exception will not be
   * re-thrown!
   */
  public void executeCommand() {
    AtomicBoolean doneSent = new AtomicBoolean(false);

    CommandResultHandler commandResultHandler = new CommandResultHandler() {
      @Override
      public void sendException(Throwable t) {
        JsonRequest.this.sendException(t);
      }

      @Override
      public void sendDone() {
        doneSent.set(true);
        JsonRequest.this.sendDone();
      }

      @Override
      public void sendData(JsonResult data) {
        try {
          String serialized = serializer.serializeWithEnvelope(requestId, JsonResultEnvelope.STATUS_DATA, data);
          synchronized (session) {
            try {
              session.getBasicRemote().sendText(serialized);
            } catch (IOException e) {
              logger.warn("Could not send data to client", e);
            }
          }
        } catch (JsonPayloadSerializerException e) {
          throw new RuntimeException("Could not serialize data", e);
        }
      }
    };

    try {
      jsonCommand.execute(ticket, commandResultHandler, commandClusterInteraction);
    } catch (AuthenticationException e) {
      try {
        String serialized =
            serializer.serializeWithEnvelope(requestId, JsonResultEnvelope.STATUS_AUTHENTICATION_EXCEPTION, null);
        session.getBasicRemote().sendText(serialized);
      } catch (IOException | JsonPayloadSerializerException e2) {
        throw new RuntimeException("Could not serialize authentication exception result.", e2);
      }
      cleanup();
      return;
    } catch (RuntimeException e) {
      sendException(e);
      logger.warn("Exception while executing command", e);
      return;
    }

    if (!AsyncJsonCommand.class.isAssignableFrom(jsonCommand.getClass()) && !doneSent.get()) {
      sendDone();
    }
  }

  /**
   * Cancel the request. This can only take effect on asynchronous commands! Otherwise this method will simply return.
   * 
   * The request will definitely be cleaned up in {@link JsonRequestRegistry}.
   */
  public void cancel() {
    if (!AsyncJsonCommand.class.isAssignableFrom(jsonCommand.getClass())) {
      cleanup();
      return;
    }

    AsyncJsonCommand asyncCommand = (AsyncJsonCommand) jsonCommand;
    asyncCommand.cancel(commandClusterInteraction);

    cleanup();
  }

  /**
   * Send a "done" to the client.
   */
  private void sendDone() {
    cleanup();

    try {
      String serialized = serializer.serializeWithEnvelope(requestId, JsonResultEnvelope.STATUS_DONE, null);
      synchronized (session) {
        try {
          session.getBasicRemote().sendText(serialized);
        } catch (IOException e) {
          logger.warn("Could not send done to client", e);
        }
      }
    } catch (JsonPayloadSerializerException e) {
      throw new RuntimeException("Could not serialize 'done'", e);
    }
  }

  /**
   * Send an exception to the client.
   * 
   * @param t
   *          The exception.
   */
  private void sendException(Throwable t) {
    cleanup();

    ExceptionJsonResult ex = new ExceptionJsonResult();
    ex.setText(t.getMessage());
    String serialized;
    try {
      serialized = serializer.serializeWithEnvelope(requestId, JsonResultEnvelope.STATUS_EXCEPTION, ex);
    } catch (JsonPayloadSerializerException e) {
      throw new RuntimeException("Could not serialize result", e);
    }
    synchronized (session) {
      try {
        session.getBasicRemote().sendText(serialized);
      } catch (IOException e) {
        logger.warn("Could not send exception to client", e);
      }
    }
  }

  private void cleanup() {
    for (Runnable r : cleanupActions)
      try {
        r.run();
      } catch (RuntimeException e) {
        logger.warn("Could not clean up request correctly", e);
        // continue with next cleanup action.
      }

    requestRegistry.unregisterRequest(session, this);
  }
}