Repository: aurora
Updated Branches:
  refs/heads/master d24092683 -> cbd8a1427


DbTaskStore perf: optimize queries scoped to a task ID.

Bugs closed: AURORA-1298

Reviewed at https://reviews.apache.org/r/35672/


Project: http://git-wip-us.apache.org/repos/asf/aurora/repo
Commit: http://git-wip-us.apache.org/repos/asf/aurora/commit/cbd8a142
Tree: http://git-wip-us.apache.org/repos/asf/aurora/tree/cbd8a142
Diff: http://git-wip-us.apache.org/repos/asf/aurora/diff/cbd8a142

Branch: refs/heads/master
Commit: cbd8a14272911285f5b08ae1f332957df8859e1e
Parents: d240926
Author: Bill Farner <[email protected]>
Authored: Mon Jun 22 09:51:36 2015 -0700
Committer: Bill Farner <[email protected]>
Committed: Mon Jun 22 09:51:36 2015 -0700

----------------------------------------------------------------------
 .../aurora/scheduler/storage/TaskStore.java     | 56 ++++++++++++++++++++
 .../scheduler/storage/db/DbTaskStore.java       | 32 +++++++----
 .../aurora/scheduler/storage/db/TaskMapper.java | 11 ++++
 .../scheduler/storage/mem/MemTaskStore.java     | 56 ++++----------------
 .../aurora/scheduler/storage/db/TaskMapper.xml  | 15 ++++--
 5 files changed, 111 insertions(+), 59 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/aurora/blob/cbd8a142/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java 
b/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java
index cceac8a..142e4ac 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/TaskStore.java
@@ -16,13 +16,18 @@ package org.apache.aurora.scheduler.storage;
 import java.util.Set;
 
 import com.google.common.base.Function;
+import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableSet;
 
+import org.apache.aurora.gen.TaskQuery;
 import org.apache.aurora.scheduler.base.Query;
+import org.apache.aurora.scheduler.base.Tasks;
 import org.apache.aurora.scheduler.storage.entities.IJobKey;
 import org.apache.aurora.scheduler.storage.entities.IScheduledTask;
 import org.apache.aurora.scheduler.storage.entities.ITaskConfig;
 
+import static com.google.common.base.CharMatcher.WHITESPACE;
+
 /**
  * Stores all tasks configured with the scheduler.
  */
@@ -106,4 +111,55 @@ public interface TaskStore {
      */
     boolean unsafeModifyInPlace(String taskId, ITaskConfig taskConfiguration);
   }
+
+  final class Util {
+    private Util() {
+      // Utility class.
+    }
+
+    public static Predicate<IScheduledTask> queryFilter(final Query.Builder 
queryBuilder) {
+      return new Predicate<IScheduledTask>() {
+        @Override
+        public boolean apply(IScheduledTask task) {
+          TaskQuery query = queryBuilder.get();
+          ITaskConfig config = task.getAssignedTask().getTask();
+          // TODO(wfarner): Investigate why blank inputs are treated specially 
for the role field.
+          if (query.getRole() != null
+              && !WHITESPACE.matchesAllOf(query.getRole())
+              && !query.getRole().equals(config.getJob().getRole())) {
+            return false;
+          }
+          if (query.getEnvironment() != null
+              && !query.getEnvironment().equals(config.getEnvironment())) {
+            return false;
+          }
+          if (query.getJobName() != null && 
!query.getJobName().equals(config.getJobName())) {
+            return false;
+          }
+
+          if (query.getJobKeysSize() > 0
+              && !query.getJobKeys().contains(config.getJob().newBuilder())) {
+            return false;
+          }
+          if (query.getTaskIds() != null && 
!query.getTaskIds().contains(Tasks.id(task))) {
+            return false;
+          }
+
+          if (query.getStatusesSize() > 0 && 
!query.getStatuses().contains(task.getStatus())) {
+            return false;
+          }
+          if (query.getSlaveHostsSize() > 0
+              && 
!query.getSlaveHosts().contains(task.getAssignedTask().getSlaveHost())) {
+            return false;
+          }
+          if (query.getInstanceIdsSize() > 0
+              && 
!query.getInstanceIds().contains(task.getAssignedTask().getInstanceId())) {
+            return false;
+          }
+
+          return true;
+        }
+      };
+    }
+  }
 }

http://git-wip-us.apache.org/repos/asf/aurora/blob/cbd8a142/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java
----------------------------------------------------------------------
diff --git 
a/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java 
b/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java
index 9b30b01..30e3469 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/DbTaskStore.java
@@ -23,6 +23,8 @@ import com.google.common.base.Function;
 import com.google.common.base.Functions;
 import com.google.common.base.Optional;
 import com.google.common.base.Preconditions;
+import com.google.common.base.Predicate;
+import com.google.common.base.Predicates;
 import com.google.common.cache.CacheBuilder;
 import com.google.common.cache.CacheLoader;
 import com.google.common.cache.LoadingCache;
@@ -291,9 +293,24 @@ class DbTaskStore implements TaskStore.Mutable {
     return Functions.compose(REPLACE_UNION_TYPES, linkPopulator);
   }
 
-  private FluentIterable<ScheduledTaskWrapper> fetchRows(Query.Builder query) {
+  private FluentIterable<IScheduledTask> matches(Query.Builder query) {
+    Iterable<ScheduledTaskWrapper> results;
+    Predicate<IScheduledTask> filter;
+    if (query.get().getTaskIdsSize() == 1) {
+      // Optimize queries that are scoped to a single task, as the dynamic SQL 
used for arbitrary
+      // queries comes with a performance penalty.
+      results = Optional.fromNullable(
+          
taskMapper.selectById(Iterables.getOnlyElement(query.get().getTaskIds())))
+          .asSet();
+      filter = Util.queryFilter(query);
+    } else {
+      results = taskMapper.select(query.get());
+      // Additional filtering is not necessary in this case, since the query 
does it for us.
+      filter = Predicates.alwaysTrue();
+    }
+
     final Function<TaskConfigRow, TaskConfig> configSaturator = 
getConfigSaturator();
-    return FluentIterable.from(taskMapper.select(query.get()))
+    return FluentIterable.from(results)
         .transform(populateAssignedPorts)
         .transform(new Function<ScheduledTaskWrapper, ScheduledTaskWrapper>() {
           @Override
@@ -304,7 +321,10 @@ class DbTaskStore implements TaskStore.Mutable {
                     task.getTask().getAssignedTask().getTask()));
             return task;
           }
-        });
+        })
+        .transform(UNWRAP)
+        .transform(IScheduledTask.FROM_BUILDER)
+        .filter(filter);
   }
 
   private final Function<ScheduledTaskWrapper, ScheduledTaskWrapper> 
populateAssignedPorts =
@@ -320,12 +340,6 @@ class DbTaskStore implements TaskStore.Mutable {
         }
       };
 
-  private FluentIterable<IScheduledTask> matches(Query.Builder query) {
-    return fetchRows(query)
-        .transform(UNWRAP)
-        .transform(IScheduledTask.FROM_BUILDER);
-  }
-
   private static final Function<ScheduledTaskWrapper, ScheduledTask> UNWRAP =
       new Function<ScheduledTaskWrapper, ScheduledTask>() {
         @Override

http://git-wip-us.apache.org/repos/asf/aurora/blob/cbd8a142/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java
----------------------------------------------------------------------
diff --git 
a/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java 
b/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java
index 8270407..7d499af 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/TaskMapper.java
@@ -16,6 +16,8 @@ package org.apache.aurora.scheduler.storage.db;
 import java.util.List;
 import java.util.Set;
 
+import javax.annotation.Nullable;
+
 import org.apache.aurora.gen.JobKey;
 import org.apache.aurora.gen.TaskQuery;
 import org.apache.aurora.scheduler.storage.db.views.AssignedPort;
@@ -44,6 +46,15 @@ interface TaskMapper {
   List<ScheduledTaskWrapper> select(TaskQuery query);
 
   /**
+   * Gets a task by ID.
+   *
+   * @param taskId ID of the task to fetch.
+   * @return Task with the specified ID.
+   */
+  @Nullable
+  ScheduledTaskWrapper selectById(@Param("taskId") String taskId);
+
+  /**
    * Gets job keys of all stored tasks.
    *
    * @return Job keys.

http://git-wip-us.apache.org/repos/asf/aurora/blob/cbd8a142/src/main/java/org/apache/aurora/scheduler/storage/mem/MemTaskStore.java
----------------------------------------------------------------------
diff --git 
a/src/main/java/org/apache/aurora/scheduler/storage/mem/MemTaskStore.java 
b/src/main/java/org/apache/aurora/scheduler/storage/mem/MemTaskStore.java
index 4b67f6b..b871575 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/mem/MemTaskStore.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/mem/MemTaskStore.java
@@ -48,7 +48,6 @@ import com.twitter.common.stats.StatsProvider;
 
 import org.apache.aurora.gen.ScheduledTask;
 import org.apache.aurora.gen.TaskConfig;
-import org.apache.aurora.gen.TaskQuery;
 import org.apache.aurora.scheduler.base.JobKeys;
 import org.apache.aurora.scheduler.base.Query;
 import org.apache.aurora.scheduler.base.Tasks;
@@ -59,8 +58,6 @@ import 
org.apache.aurora.scheduler.storage.entities.ITaskConfig;
 
 import static java.util.Objects.requireNonNull;
 
-import static com.google.common.base.CharMatcher.WHITESPACE;
-
 /**
  * An in-memory task store.
  */
@@ -247,48 +244,15 @@ class MemTaskStore implements TaskStore.Mutable {
     }
   }
 
-  private static Predicate<Task> queryFilter(final TaskQuery query) {
-    return new Predicate<Task>() {
-      @Override
-      public boolean apply(Task canonicalTask) {
-        IScheduledTask task = canonicalTask.storedTask;
-        ITaskConfig config = task.getAssignedTask().getTask();
-        if (query.getRole() != null
-            && !WHITESPACE.matchesAllOf(query.getRole())
-            && !query.getRole().equals(config.getJob().getRole())) {
-          return false;
-        }
-        if (query.getEnvironment() != null
-            && !query.getEnvironment().equals(config.getEnvironment())) {
-          return false;
-        }
-        if (query.getJobName() != null && 
!query.getJobName().equals(config.getJobName())) {
-          return false;
-        }
-
-        if (query.getJobKeysSize() > 0
-            && !query.getJobKeys().contains(config.getJob().newBuilder())) {
-          return false;
-        }
-        if (query.getTaskIds() != null && 
!query.getTaskIds().contains(Tasks.id(task))) {
-            return false;
-        }
-
-        if (query.getStatusesSize() > 0 && 
!query.getStatuses().contains(task.getStatus())) {
-          return false;
-        }
-        if (query.getSlaveHostsSize() > 0
-            && 
!query.getSlaveHosts().contains(task.getAssignedTask().getSlaveHost())) {
-          return false;
-        }
-        if (query.getInstanceIdsSize() > 0
-            && 
!query.getInstanceIds().contains(task.getAssignedTask().getInstanceId())) {
-          return false;
-        }
-
-        return true;
-      }
-    };
+  private static Predicate<Task> queryFilter(Query.Builder query) {
+    return Predicates.compose(
+        Util.queryFilter(query),
+        new Function<Task, IScheduledTask>() {
+          @Override
+          public IScheduledTask apply(Task canonicalTask) {
+            return canonicalTask.storedTask;
+          }
+        });
   }
 
   private Iterable<Task> fromIdIndex(Iterable<String> taskIds) {
@@ -323,7 +287,7 @@ class MemTaskStore implements TaskStore.Mutable {
       }
     }
 
-    return FluentIterable.from(from.get()).filter(queryFilter(query.get()));
+    return FluentIterable.from(from.get()).filter(queryFilter(query));
   }
 
   private static final Function<Task, IScheduledTask> TO_SCHEDULED =

http://git-wip-us.apache.org/repos/asf/aurora/blob/cbd8a142/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml
----------------------------------------------------------------------
diff --git 
a/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml 
b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml
index 7c27f37..e9660b7 100644
--- a/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml
+++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/TaskMapper.xml
@@ -89,10 +89,7 @@
     <association property="task" resultMap="scheduledTaskMap"/>
   </resultMap>
 
-  <!-- TODO(wfarner): Consider adding selectById and/or selectByIds methods.  
The dynamic SQL here
-       seems to come with a ~3x performance hit.
-   -->
-  <select id="select" resultMap="taskWrapperMap">
+  <sql id="unscopedSelect">
     SELECT
       t.id AS row_id,
       t.task_config_row_id AS task_config_row_id,
@@ -110,6 +107,16 @@
     INNER JOIN task_configs as c ON c.id = t.task_config_row_id
     INNER JOIN job_keys AS j ON j.id = c.job_key_id
     LEFT OUTER JOIN host_attributes AS h ON h.id = t.slave_row_id
+  </sql>
+
+  <select id="selectById" resultMap="taskWrapperMap">
+    <include refid="unscopedSelect"/>
+    WHERE
+      t.task_id = #{taskId}
+  </select>
+
+  <select id="select" resultMap="taskWrapperMap">
+    <include refid="unscopedSelect"/>
     <where>
       <if test="role != null">
         j.role = #{role}

Reply via email to