DistributedResourcePoolTest.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.zeppelin.resource;

import com.google.gson.Gson;
import org.apache.zeppelin.user.AuthenticationInfo;
import org.apache.zeppelin.display.GUI;
import org.apache.zeppelin.interpreter.*;
import org.apache.zeppelin.interpreter.remote.RemoteInterpreter;
import org.apache.zeppelin.interpreter.remote.RemoteInterpreterEventPoller;
import org.apache.zeppelin.interpreter.remote.mock.MockInterpreterResourcePool;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;

import java.io.File;
import java.util.HashMap;
import java.util.LinkedList;
import java.util.Properties;

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

/**
 * Unittest for DistributedResourcePool
 */
public class DistributedResourcePoolTest {
  private static final String INTERPRETER_SCRIPT =
          System.getProperty("os.name").startsWith("Windows") ?
                  "../bin/interpreter.cmd" :
                  "../bin/interpreter.sh";
  private InterpreterGroup intpGroup1;
  private InterpreterGroup intpGroup2;
  private HashMap<String, String> env;
  private RemoteInterpreter intp1;
  private RemoteInterpreter intp2;
  private InterpreterContext context;
  private RemoteInterpreterEventPoller eventPoller1;
  private RemoteInterpreterEventPoller eventPoller2;


  @Before
  public void setUp() throws Exception {
    env = new HashMap<>();
    env.put("ZEPPELIN_CLASSPATH", new File("./target/test-classes").getAbsolutePath());

    Properties p = new Properties();

    intp1 = new RemoteInterpreter(
        p,
        "note",
        MockInterpreterResourcePool.class.getName(),
        new File(INTERPRETER_SCRIPT).getAbsolutePath(),
        "fake",
        "fakeRepo",
        env,
        10 * 1000,
        null,
        null,
        "anonymous",
        false
    );

    intpGroup1 = new InterpreterGroup("intpGroup1");
    intpGroup1.put("note", new LinkedList<Interpreter>());
    intpGroup1.get("note").add(intp1);
    intp1.setInterpreterGroup(intpGroup1);

    intp2 = new RemoteInterpreter(
        p,
        "note",
        MockInterpreterResourcePool.class.getName(),
        new File(INTERPRETER_SCRIPT).getAbsolutePath(),
        "fake",
        "fakeRepo",        
        env,
        10 * 1000,
        null,
        null,
        "anonymous",
        false
    );

    intpGroup2 = new InterpreterGroup("intpGroup2");
    intpGroup2.put("note", new LinkedList<Interpreter>());
    intpGroup2.get("note").add(intp2);
    intp2.setInterpreterGroup(intpGroup2);

    context = new InterpreterContext(
        "note",
        "id",
        null,
        "title",
        "text",
        new AuthenticationInfo(),
        new HashMap<String, Object>(),
        new GUI(),
        null,
        null,
        new LinkedList<InterpreterContextRunner>(),
        null);

    intp1.open();
    intp2.open();

    eventPoller1 = new RemoteInterpreterEventPoller(null, null);
    eventPoller1.setInterpreterGroup(intpGroup1);
    eventPoller1.setInterpreterProcess(intpGroup1.getRemoteInterpreterProcess());

    eventPoller2 = new RemoteInterpreterEventPoller(null, null);
    eventPoller2.setInterpreterGroup(intpGroup2);
    eventPoller2.setInterpreterProcess(intpGroup2.getRemoteInterpreterProcess());

    eventPoller1.start();
    eventPoller2.start();
  }

  @After
  public void tearDown() throws Exception {
    eventPoller1.shutdown();
    intp1.close();
    intpGroup1.close();
    eventPoller2.shutdown();
    intp2.close();
    intpGroup2.close();
  }

  @Test
  public void testRemoteDistributedResourcePool() {
    Gson gson = new Gson();
    InterpreterResult ret;
    intp1.interpret("put key1 value1", context);
    intp2.interpret("put key2 value2", context);

    ret = intp1.interpret("getAll", context);
    assertEquals(2, gson.fromJson(ret.message().get(0).getData(), ResourceSet.class).size());

    ret = intp2.interpret("getAll", context);
    assertEquals(2, gson.fromJson(ret.message().get(0).getData(), ResourceSet.class).size());

    ret = intp1.interpret("get key1", context);
    assertEquals("value1", gson.fromJson(ret.message().get(0).getData(), String.class));

    ret = intp1.interpret("get key2", context);
    assertEquals("value2", gson.fromJson(ret.message().get(0).getData(), String.class));
  }

  @Test
  public void testDistributedResourcePool() {
    final LocalResourcePool pool2 = new LocalResourcePool("pool2");
    final LocalResourcePool pool3 = new LocalResourcePool("pool3");

    DistributedResourcePool pool1 = new DistributedResourcePool("pool1", new ResourcePoolConnector() {
      @Override
      public ResourceSet getAllResources() {
        ResourceSet set = pool2.getAll();
        set.addAll(pool3.getAll());

        ResourceSet remoteSet = new ResourceSet();
        Gson gson = new Gson();
        for (Resource s : set) {
          RemoteResource remoteResource = gson.fromJson(gson.toJson(s), RemoteResource.class);
          remoteResource.setResourcePoolConnector(this);
          remoteSet.add(remoteResource);
        }
        return remoteSet;
      }

      @Override
      public Object readResource(ResourceId id) {
        if (id.getResourcePoolId().equals(pool2.id())) {
          return pool2.get(id.getName()).get();
        }
        if (id.getResourcePoolId().equals(pool3.id())) {
          return pool3.get(id.getName()).get();
        }
        return null;
      }

      @Override
      public Object invokeMethod(ResourceId id, String methodName, Class[] paramTypes, Object[] params) {
        return null;
      }

      @Override
      public Resource invokeMethod(ResourceId id, String methodName, Class[] paramTypes, Object[]
          params, String returnResourceName) {
        return null;
      }
    });

    assertEquals(0, pool1.getAll().size());


    // test get() can get from pool
    pool2.put("object1", "value2");
    assertEquals(1, pool1.getAll().size());
    assertTrue(pool1.get("object1").isRemote());
    assertEquals("value2", pool1.get("object1").get());

    // test get() is locality aware
    pool1.put("object1", "value1");
    assertEquals(1, pool2.getAll().size());
    assertEquals("value1", pool1.get("object1").get());

    // test getAll() is locality aware
    assertEquals("value1", pool1.getAll().get(0).get());
    assertEquals("value2", pool1.getAll().get(1).get());
  }

  @Test
  public void testResourcePoolUtils() {
    Gson gson = new Gson();
    InterpreterResult ret;

    // when create some resources
    intp1.interpret("put note1:paragraph1:key1 value1", context);
    intp1.interpret("put note1:paragraph2:key1 value2", context);
    intp2.interpret("put note2:paragraph1:key1 value1", context);
    intp2.interpret("put note2:paragraph2:key2 value2", context);


    // then get all resources.
    assertEquals(4, ResourcePoolUtils.getAllResources().size());

    // when remove all resources from note1
    ResourcePoolUtils.removeResourcesBelongsToNote("note1");

    // then resources should be removed.
    assertEquals(2, ResourcePoolUtils.getAllResources().size());
    assertEquals("", gson.fromJson(
        intp1.interpret("get note1:paragraph1:key1", context).message().get(0).getData(),
        String.class));
    assertEquals("", gson.fromJson(
        intp1.interpret("get note1:paragraph2:key1", context).message().get(0).getData(),
        String.class));


    // when remove all resources from note2:paragraph1
    ResourcePoolUtils.removeResourcesBelongsToParagraph("note2", "paragraph1");

    // then 1
    assertEquals(1, ResourcePoolUtils.getAllResources().size());
    assertEquals("value2", gson.fromJson(
        intp1.interpret("get note2:paragraph2:key2", context).message().get(0).getData(),
        String.class));

  }

  @Test
  public void testResourceInvokeMethod() {
    Gson gson = new Gson();
    InterpreterResult ret;
    intp1.interpret("put key1 hey", context);
    intp2.interpret("put key2 world", context);

    // invoke method in local resource pool
    ret = intp1.interpret("invoke key1 length", context);
    assertEquals("3", ret.message().get(0).getData());

    // invoke method in remote resource pool
    ret = intp1.interpret("invoke key2 length", context);
    assertEquals("5", ret.message().get(0).getData());

    // make sure no resources are automatically created
    ret = intp1.interpret("getAll", context);
    assertEquals(2, gson.fromJson(ret.message().get(0).getData(), ResourceSet.class).size());

    // invoke method in local resource pool and save result
    ret = intp1.interpret("invoke key1 length ret1", context);
    assertEquals("3", ret.message().get(0).getData());

    ret = intp1.interpret("getAll", context);
    assertEquals(3, gson.fromJson(ret.message().get(0).getData(), ResourceSet.class).size());

    ret = intp1.interpret("get ret1", context);
    assertEquals("3", gson.fromJson(ret.message().get(0).getData(), String.class));

    // invoke method in remote resource pool and save result
    ret = intp1.interpret("invoke key2 length ret2", context);
    assertEquals("5", ret.message().get(0).getData());

    ret = intp1.interpret("getAll", context);
    assertEquals(4, gson.fromJson(ret.message().get(0).getData(), ResourceSet.class).size());

    ret = intp1.interpret("get ret2", context);
    assertEquals("5", gson.fromJson(ret.message().get(0).getData(), String.class));
  }
}