giraph-block-app/src/test/java/org/apache/giraph/block_app/framework/BlockApiHandleTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.giraph.block_app.framework;

import org.apache.giraph.block_app.framework.api.BlockApiHandle;
import org.apache.giraph.block_app.framework.api.BlockMasterApi;
import org.apache.giraph.block_app.framework.api.BlockWorkerContextReceiveApi;
import org.apache.giraph.block_app.framework.api.BlockWorkerContextSendApi;
import org.apache.giraph.block_app.framework.api.BlockWorkerReceiveApi;
import org.apache.giraph.block_app.framework.api.BlockWorkerSendApi;
import org.apache.giraph.block_app.framework.api.local.LocalBlockRunner;
import org.apache.giraph.block_app.framework.block.Block;
import org.apache.giraph.block_app.framework.block.BlockWithApiHandle;
import org.apache.giraph.block_app.framework.block.PieceCount;
import org.apache.giraph.block_app.framework.piece.AbstractPiece;
import org.apache.giraph.block_app.framework.piece.DefaultParentPiece;
import org.apache.giraph.block_app.framework.piece.interfaces.VertexReceiver;
import org.apache.giraph.block_app.framework.piece.interfaces.VertexSender;
import org.apache.giraph.conf.GiraphConfiguration;
import org.apache.giraph.conf.GiraphConstants;
import org.apache.giraph.function.Consumer;
import org.apache.giraph.graph.Vertex;
import org.apache.giraph.utils.TestGraph;
import org.apache.hadoop.io.DoubleWritable;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.io.NullWritable;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableComparable;
import org.junit.Test;

import java.util.Iterator;
import java.util.List;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;

/**
 * Test the use of {@link BlockApiHandle}.
 */
public class BlockApiHandleTest {

  private static GiraphConfiguration createConf() {
    GiraphConfiguration conf = new GiraphConfiguration();
    GiraphConstants.VERTEX_ID_CLASS.set(conf, LongWritable.class);
    GiraphConstants.VERTEX_VALUE_CLASS.set(conf, LongWritable.class);
    GiraphConstants.EDGE_VALUE_CLASS.set(conf, NullWritable.class);
    return conf;
  }

  private static TestGraph<LongWritable, LongWritable, NullWritable>
  createTestGraph() {
    TestGraph<LongWritable, LongWritable, NullWritable> graph =
      new TestGraph<>(createConf());
    graph.addVertex(new LongWritable(1), new LongWritable());
    graph.addVertex(new LongWritable(2), new LongWritable());
    graph.addVertex(new LongWritable(3), new LongWritable());
    graph.addVertex(new LongWritable(4), new LongWritable());
    graph.addEdge(new LongWritable(1), new LongWritable(2), NullWritable.get());
    graph.addEdge(new LongWritable(2), new LongWritable(1), NullWritable.get());
    graph.addEdge(new LongWritable(2), new LongWritable(3), NullWritable.get());
    graph.addEdge(new LongWritable(3), new LongWritable(2), NullWritable.get());
    return graph;
  }

  public static class DummyObjectWithApiHandle {

    private BlockApiHandle handle;

    public DummyObjectWithApiHandle(BlockApiHandle handle) {
      this.handle = handle;
    }

    public void doSomethingAtWorker() {
      // checking that the handles have been set
      assertFalse(handle.isMasterApiSet());
      assertFalse(handle.isWorkerContextReceiveApiSet());
      assertFalse(handle.isWorkerContextSendApiSet());
      assertEquals(1, handle.getWorkerReceiveApi().getWorkerCount());
      assertEquals(0, handle.getWorkerSendApi().getMyWorkerIndex());
    }

    public void doSomethingAtWorkerContext() {
      // checking that the handles have been set
      assertFalse(handle.isMasterApiSet());
      assertFalse(handle.isWorkerReceiveApiSet());
      assertFalse(handle.isWorkerSendApiSet());
      assertEquals(1, handle.getWorkerContextReceiveApi().getWorkerCount());
      assertEquals(0, handle.getWorkerContextSendApi().getMyWorkerIndex());
    }

    public void doSomethingAtMaster() {
      // checking that the handles have been set
      assertEquals(1, handle.getMasterApi().getWorkerCount());
      assertFalse(handle.isWorkerReceiveApiSet());
      assertFalse(handle.isWorkerSendApiSet());
      assertFalse(handle.isWorkerContextReceiveApiSet());
      assertFalse(handle.isWorkerContextSendApiSet());
    }
  }

  public static class TestPiece extends DefaultParentPiece<WritableComparable,
      LongWritable, Writable, NullWritable, Object, DoubleWritable, Object> {

    private DummyObjectWithApiHandle object;

    public TestPiece(DummyObjectWithApiHandle object) {
      this.object = object;
    }

    @Override
    public VertexSender<WritableComparable, LongWritable, Writable>
    getVertexSender(final BlockWorkerSendApi<WritableComparable, LongWritable,
      Writable, NullWritable> workerApi, Object executionStage) {
      return new InnerVertexSender() {
        @Override
        public void vertexSend(
          Vertex<WritableComparable, LongWritable, Writable> vertex) {
          object.doSomethingAtWorker();
        }
      };
    }

    @Override
    public VertexReceiver<WritableComparable, LongWritable, Writable,
      NullWritable> getVertexReceiver(
      BlockWorkerReceiveApi<WritableComparable> workerApi,
      Object executionStage) {
      return new InnerVertexReceiver() {
        @Override
        public void vertexReceive(
          Vertex<WritableComparable, LongWritable, Writable> vertex,
          Iterable<NullWritable> messages) {
          object.doSomethingAtWorker();
        }
      };
    }

    public void workerContextSend(BlockWorkerContextSendApi<WritableComparable,
        DoubleWritable> workerContextApi, Object executionStage,
        Writable workerValue) {
      object.doSomethingAtWorkerContext();
    }

    /**
     * Override to have worker context receive computation.
     *
     * Called once per worker, before all vertices are going to be processed
     * with getVertexReceiver.
     */
    public void workerContextReceive(
      BlockWorkerContextReceiveApi workerContextApi, Object executionStage,
      Object workerValue, List<DoubleWritable> workerMessages) {
      object.doSomethingAtWorkerContext();
    }

    @Override
    public void masterCompute(BlockMasterApi masterApi, Object executionStage) {
      object.doSomethingAtMaster();
    }

    @Override
    protected Class<NullWritable> getMessageClass() {
      return NullWritable.class;
    }
  }

  @Test
  public void testBlockApiHandle() {
    TestGraph<LongWritable, LongWritable, NullWritable> graph =
      createTestGraph();

    final BlockApiHandle handle = new BlockApiHandle();
    final DefaultParentPiece piece =
      new TestPiece(new DummyObjectWithApiHandle(handle));

    Block block = new BlockWithApiHandle() {
      @Override
      public Iterator<AbstractPiece> iterator() {
        return piece.iterator();
      }

      @Override
      public void forAllPossiblePieces(Consumer<AbstractPiece> consumer) {
        piece.forAllPossiblePieces(consumer);
      }

      @Override
      public PieceCount getPieceCount() {
        return piece.getPieceCount();
      }

      @Override
      public BlockApiHandle getBlockApiHandle() {
        return handle;
      }
    };

    BlockUtils.BLOCK_WORKER_CONTEXT_VALUE_CLASS.set(
      graph.getConf(), Object.class);
    LocalBlockRunner.runBlock(graph, block, new Object());
  }
}