SAMZA-1552: Host affinity improvements - Improve matching of hosts to allocated resources

Author: Jagadish <[email protected]>

Reviewers: Prateek <[email protected]>

Closes apache#401 from vjagadish1989/host-affinity-fix

Author: jagadish-v0

samza-core/src/main/java/org/apache/samza/clustermanager/ResourceRequestState.java

@@ -171,7 +171,9 @@ private void addToAllocatedResourceList(String host, SamzaResource samzaResource
   public void updateStateAfterAssignment(SamzaResourceRequest request, String assignedHost, SamzaResource samzaResource) {
     synchronized (lock) {
       requestsQueue.remove(request);
+      // A reference for the resource could either be held in the preferred host buffer or in the ANY_HOST buffer.
       allocatedResources.get(assignedHost).remove(samzaResource);
+      allocatedResources.get(ANY_HOST).remove(samzaResource);
       if (hostAffinityEnabled) {
         // assignedHost may not always be the preferred host.
         // Hence, we should safely decrement the counter for the preferredHost
@@ -266,18 +268,33 @@ private List<String> getAllocatedHosts() {
   /**
    * Retrieves, but does not remove, the first allocated resource on the specified host.
    *
-   * @param host  the host for which a resource is needed.
-   * @return      the first {@link SamzaResource} allocated for the specified host or {@code null} if there isn't one.
+   * @param host the host for which a resource is needed.
+   * @return a {@link SamzaResource} allocated for the specified host or {@code null} if there isn't one.
    */
-
   public SamzaResource peekResource(String host) {
     synchronized (lock) {
-      List<SamzaResource> resourcesOnTheHost = this.allocatedResources.get(host);
+      List<SamzaResource> resourcesOnPreferredHostBuffer = this.allocatedResources.get(host);
+      List<SamzaResource> resourcesOnAnyHostBuffer = this.allocatedResources.get(ANY_HOST);
 
-      if (resourcesOnTheHost == null || resourcesOnTheHost.isEmpty()) {
+      // First search for the preferred host buffers
+      if (resourcesOnPreferredHostBuffer != null && !resourcesOnPreferredHostBuffer.isEmpty()) {
+        SamzaResource resource = resourcesOnPreferredHostBuffer.get(0);
+        log.info("Returning a buffered resource: {} for {} from preferred-host buffer.", resource.getResourceID(), host);
+        return resource;
+      } else if (resourcesOnAnyHostBuffer != null && !resourcesOnAnyHostBuffer.isEmpty()) {
+        // If preferred host buffers are empty, scan the ANY_HOST buffer
+        log.debug("No resources on preferred-host buffer. Scanning ANY_HOST buffer");
+        SamzaResource resource = resourcesOnAnyHostBuffer.stream()
+            .filter(resrc -> resrc.getHost().equals(host))
+            .findAny().orElse(null);
+        if (resource != null) {
+          log.info("Returning a buffered resource: {} for {} from ANY_HOST buffer.", resource.getResourceID(), host);
+        }
+        return resource;
+      } else {
+        log.debug("Cannot find any resource in the ANY_HOST buffer for {} because both buffers are empty", host);
         return null;
       }
-      return resourcesOnTheHost.get(0);
     }
   }
 

samza-core/src/test/java/org/apache/samza/clustermanager/MockHostAwareContainerAllocator.java

@@ -0,0 +1,68 @@
+/*
+ * 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.samza.clustermanager;
+
+import org.apache.samza.config.Config;
+
+import java.lang.reflect.Field;
+import java.util.Map;
+import java.util.concurrent.Semaphore;
+import java.util.concurrent.TimeUnit;
+
+public class MockHostAwareContainerAllocator extends HostAwareContainerAllocator {
+  private static final int ALLOCATOR_TIMEOUT_MS = 10000;
+  private Semaphore semaphore = new Semaphore(0);
+
+  public MockHostAwareContainerAllocator(ClusterResourceManager manager,
+      Config config, SamzaApplicationState state) {
+    super(manager, ALLOCATOR_TIMEOUT_MS, config, state);
+  }
+
+  /**
+   * Causes the current thread to block until the expected number of containers have started.
+   *
+   * @param numExpectedContainers the number of containers expected to start
+   * @param timeout the maximum time to wait
+   * @param unit the time unit of the {@code timeout} argument
+   *
+   * @return a boolean that specifies whether containers started within the timeout.
+   * @throws InterruptedException  if the current thread is interrupted while waiting
+   */
+  boolean awaitContainersStart(int numExpectedContainers, long timeout, TimeUnit unit) throws InterruptedException {
+    return semaphore.tryAcquire(numExpectedContainers, timeout, unit);
+  }
+
+  @Override
+  public void requestResources(Map<String, String> containerToHostMappings) {
+    super.requestResources(containerToHostMappings);
+  }
+
+  public ResourceRequestState getContainerRequestState() throws Exception {
+    Field field = AbstractContainerAllocator.class.getDeclaredField("resourceRequestState");
+    field.setAccessible(true);
+
+    return (ResourceRequestState) field.get(this);
+  }
+
+  @Override
+  protected void runStreamProcessor(SamzaResourceRequest request, String preferredHost) {
+    super.runStreamProcessor(request, preferredHost);
+    semaphore.release();
+  }
+}