aurora-commits mailing list archives

Site index · List index
Message view « Date » · « Thread »
Top « Date » · « Thread »
From kevi...@apache.org
Subject git commit: Initial attempt at h2/DB storage implementation (LockStore only)
Date Wed, 28 May 2014 01:05:39 GMT
Repository: incubator-aurora
Updated Branches:
  refs/heads/master 2753763e7 -> 1b6a55e46


Initial attempt at h2/DB storage implementation (LockStore only)

Lays down the groundwork for h2/MyBatis integration in the
aurora scheduler, and creates a bridge so that the existing
in-memory stores can co-exist with the new db stores. This
allows us to incrementally replace each storage implementation in
org.apache.aurora.scheduler.storage.mem. In this patch I have replaced
MemLockStore with DbLockStore.

Testing Done:
./gradlew clean build

Bugs closed: AURORA-335

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


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

Branch: refs/heads/master
Commit: 1b6a55e4619b6ae9e016b626d30c9aec1a189e64
Parents: 2753763
Author: David McLaughlin <david@dmclaughlin.com>
Authored: Tue May 27 17:56:53 2014 -0700
Committer: Kevin Sweeney <kevints@apache.org>
Committed: Tue May 27 18:05:10 2014 -0700

----------------------------------------------------------------------
 build.gradle                                    |   3 +
 .../aurora/scheduler/app/SchedulerMain.java     |   8 +
 .../apache/aurora/scheduler/base/JobKeys.java   |   1 +
 .../aurora/scheduler/storage/Storage.java       |  16 +-
 .../scheduler/storage/db/DbLockStore.java       |  93 +++++++++
 .../aurora/scheduler/storage/db/DbModule.java   | 107 ++++++++++
 .../aurora/scheduler/storage/db/DbStorage.java  | 162 +++++++++++++++
 .../scheduler/storage/db/JobKeyMapper.java      |  40 ++++
 .../scheduler/storage/db/LockKeyMapper.java     |  49 +++++
 .../aurora/scheduler/storage/db/LockMapper.java |  53 +++++
 .../scheduler/storage/db/MigrationModule.java   |  48 +++++
 .../scheduler/storage/db/views/LockRow.java     |  46 +++++
 .../scheduler/storage/log/LogStorage.java       |   1 +
 .../scheduler/storage/log/LogStorageModule.java |   2 +-
 .../scheduler/storage/mem/MemLockStore.java     |  58 ------
 .../scheduler/storage/mem/MemStorage.java       |  45 +++-
 .../scheduler/storage/mem/MemStorageModule.java |   2 -
 .../scheduler/storage/db/JobKeyMapper.xml       |  33 +++
 .../aurora/scheduler/storage/db/LockMapper.xml  |  69 +++++++
 .../aurora/scheduler/storage/db/schema.sql      |  21 ++
 .../scheduler/storage/db/DbLockStoreTest.java   | 203 +++++++++++++++++++
 .../scheduler/storage/log/LogStorageTest.java   |   2 +
 .../scheduler/storage/mem/MemLockStoreTest.java |  77 -------
 23 files changed, 983 insertions(+), 156 deletions(-)
----------------------------------------------------------------------


http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/build.gradle
----------------------------------------------------------------------
diff --git a/build.gradle b/build.gradle
index 646cdd2..9dea858 100644
--- a/build.gradle
+++ b/build.gradle
@@ -152,6 +152,7 @@ dependencies {
   compile guavaDep
   compile 'com.google.inject:guice:3.0'
   compile 'com.google.protobuf:protobuf-java:2.5.0'
+  compile 'com.h2database:h2:1.4.177'
   compile "com.sun.jersey:jersey-core:${jerseyRev}"
   compile "com.sun.jersey:jersey-json:${jerseyRev}"
   compile "com.sun.jersey:jersey-server:${jerseyRev}"
@@ -165,6 +166,8 @@ dependencies {
   compile 'org.apache.mesos:mesos:0.18.0'
   compile thriftLib
   compile 'org.apache.zookeeper:zookeeper:3.3.4'
+  compile 'org.mybatis:mybatis:3.2.7'
+  compile 'org.mybatis:mybatis-guice:3.6'
   compile 'org.quartz-scheduler:quartz:2.2.1'
   compile "org.slf4j:slf4j-api:${slf4jRev}"
   compile "org.slf4j:slf4j-jdk14:${slf4jRev}"

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/app/SchedulerMain.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/app/SchedulerMain.java b/src/main/java/org/apache/aurora/scheduler/app/SchedulerMain.java
index 80252bd..321ac3a 100644
--- a/src/main/java/org/apache/aurora/scheduler/app/SchedulerMain.java
+++ b/src/main/java/org/apache/aurora/scheduler/app/SchedulerMain.java
@@ -61,9 +61,12 @@ import org.apache.aurora.scheduler.cron.quartz.CronModule;
 import org.apache.aurora.scheduler.local.IsolatedSchedulerModule;
 import org.apache.aurora.scheduler.log.mesos.MesosLogStreamModule;
 import org.apache.aurora.scheduler.storage.backup.BackupModule;
+import org.apache.aurora.scheduler.storage.db.DbModule;
+import org.apache.aurora.scheduler.storage.db.MigrationModule;
 import org.apache.aurora.scheduler.storage.log.LogStorage;
 import org.apache.aurora.scheduler.storage.log.LogStorageModule;
 import org.apache.aurora.scheduler.storage.log.SnapshotStoreImpl;
+import org.apache.aurora.scheduler.storage.mem.MemStorage.Delegated;
 import org.apache.aurora.scheduler.storage.mem.MemStorageModule;
 import org.apache.aurora.scheduler.thrift.ThriftConfiguration;
 import org.apache.aurora.scheduler.thrift.ThriftModule;
@@ -159,6 +162,11 @@ public class SchedulerMain extends AbstractApplication {
         .add(new LogStorageModule())
         .add(new MemStorageModule(Bindings.annotatedKeyFactory(LogStorage.WriteBehind.class)))
         .add(new CronModule())
+        .add(new DbModule(Bindings.annotatedKeyFactory(Delegated.class)))
+        .add(new MigrationModule(
+            Bindings.annotatedKeyFactory(LogStorage.WriteBehind.class),
+            Bindings.annotatedKeyFactory(Delegated.class))
+        )
         .add(new ThriftModule())
         .add(new ThriftAuthModule())
         .build();

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/base/JobKeys.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/base/JobKeys.java b/src/main/java/org/apache/aurora/scheduler/base/JobKeys.java
index ec53232..a76c3fa 100644
--- a/src/main/java/org/apache/aurora/scheduler/base/JobKeys.java
+++ b/src/main/java/org/apache/aurora/scheduler/base/JobKeys.java
@@ -15,6 +15,7 @@ package org.apache.aurora.scheduler.base;
 
 import java.util.List;
 import java.util.Set;
+
 import javax.annotation.Nullable;
 
 import com.google.common.base.Function;

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/Storage.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/Storage.java b/src/main/java/org/apache/aurora/scheduler/storage/Storage.java
index bbbd7dc..baa5ab6 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/Storage.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/Storage.java
@@ -206,18 +206,18 @@ public interface Storage {
   <T, E extends Exception> T write(MutateWork<T, E> work) throws StorageException, E;
 
   /**
+   * Requests the underlying storage prepare its data set; ie: initialize schemas, begin syncing
+   * out of date data, etc.  This method should not block.
+   *
+   * @throws StorageException if there was a problem preparing storage.
+   */
+   void prepare() throws StorageException;
+
+  /**
    * A non-volatile storage that has additional methods to control its lifecycle.
    */
   interface NonVolatileStorage extends Storage {
     /**
-     * Requests the underlying storage prepare its data set; ie: initialize schemas, begin syncing
-     * out of date data, etc.  This method should not block.
-     *
-     * @throws StorageException if there was a problem preparing storage.
-     */
-    void prepare() throws StorageException;
-
-    /**
      * Prepares the underlying storage for serving traffic.
      *
      * @param initializationLogic work to perform after this storage system is ready but before

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/db/DbLockStore.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/DbLockStore.java b/src/main/java/org/apache/aurora/scheduler/storage/db/DbLockStore.java
new file mode 100644
index 0000000..0e7f52c
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/DbLockStore.java
@@ -0,0 +1,93 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db;
+
+import java.util.Set;
+import java.util.logging.Logger;
+
+import com.google.common.base.Function;
+import com.google.common.base.Optional;
+import com.google.common.collect.FluentIterable;
+import com.google.inject.Inject;
+
+import org.apache.aurora.scheduler.storage.LockStore;
+import org.apache.aurora.scheduler.storage.db.views.LockRow;
+import org.apache.aurora.scheduler.storage.entities.ILock;
+import org.apache.aurora.scheduler.storage.entities.ILockKey;
+import org.apache.ibatis.exceptions.PersistenceException;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+/**
+ *  A relational database-backed lock store.
+ */
+class DbLockStore implements LockStore.Mutable {
+
+  private static final Logger LOG = Logger.getLogger(DbLockStore.class.getName());
+
+  private final LockMapper mapper;
+  private final LockKeyMapper lockKeyMapper;
+
+  @Inject
+  DbLockStore(LockMapper mapper, LockKeyMapper lockKeyMapper) {
+    this.mapper = checkNotNull(mapper);
+    this.lockKeyMapper = checkNotNull(lockKeyMapper);
+  }
+
+  @Override
+  public void saveLock(ILock lock) {
+    try {
+      lockKeyMapper.insert(lock.getKey().newBuilder());
+    } catch (PersistenceException e) {
+      LOG.fine("DB write error for key: " + lock.getKey());
+      // TODO(davmclau): We purposely swallow duplicate key exceptions here
+      // but we should verify it _is_ a duplicate key error so we can
+      // give better logging for unexpected errors. That is
+      // made tricky by this: https://code.google.com/p/mybatis/issues/detail?id=249
+      // It is currently harmless to let this fall through, as if the
+      // write failed and key doesn't exist, the next write will fail anyway.
+    }
+    mapper.insert(lock.newBuilder());
+  }
+
+  @Override
+  public void removeLock(ILockKey lockKey) {
+    mapper.delete(lockKey.newBuilder());
+  }
+
+  @Override
+  public void deleteLocks() {
+    mapper.truncate();
+  }
+
+  /**
+   * LockRow converter to satisfy the ILock interface.
+   */
+  private static final Function<LockRow, ILock> TO_ROW = new Function<LockRow, ILock>() {
+    @Override
+    public ILock apply(LockRow input) {
+      return ILock.build(input.getLock());
+    }
+  };
+
+  @Override
+  public Set<ILock> fetchLocks() {
+    return FluentIterable.from(mapper.selectAll()).transform(TO_ROW).toSet();
+  }
+
+  @Override
+  public Optional<ILock> fetchLock(ILockKey lockKey) {
+    return Optional.fromNullable(mapper.select(lockKey.newBuilder())).transform(TO_ROW);
+  }
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/db/DbModule.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/DbModule.java b/src/main/java/org/apache/aurora/scheduler/storage/db/DbModule.java
new file mode 100644
index 0000000..e365cd6
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/DbModule.java
@@ -0,0 +1,107 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db;
+
+import java.util.Properties;
+
+import javax.inject.Singleton;
+
+import com.google.inject.Key;
+import com.google.inject.PrivateModule;
+import com.google.inject.name.Names;
+import com.twitter.common.inject.Bindings;
+
+import org.apache.aurora.scheduler.storage.LockStore;
+import org.apache.aurora.scheduler.storage.Storage;
+import org.apache.ibatis.session.AutoMappingBehavior;
+import org.apache.ibatis.session.SqlSessionFactory;
+import org.apache.ibatis.transaction.jdbc.JdbcTransactionFactory;
+import org.h2.Driver;
+import org.mybatis.guice.MyBatisModule;
+import org.mybatis.guice.datasource.builtin.PooledDataSourceProvider;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+import static com.google.inject.name.Names.named;
+
+/**
+ * Binding module for a relational database storage system.
+ * <p>
+ *   Currently only exposes bindings for:
+ *   <ul>
+ *     <li>{@link org.apache.aurora.scheduler.storage.db.DbStorage}</li>
+ *     <li>{@link org.apache.ibatis.session.SqlSessionFactory}</li>
+ *     <li>Keys provided by the provided{@code keyFactory} for:
+ *        <ul>
+ *          <li>{@link org.apache.aurora.scheduler.storage.LockStore.Mutable}</li>
+ *        </ul>
+ *     </li>
+ *   </ul>
+ * </p>
+ */
+public class DbModule extends PrivateModule {
+
+  private final Bindings.KeyFactory keyFactory;
+
+  public DbModule(Bindings.KeyFactory keyFactory) {
+    this.keyFactory = checkNotNull(keyFactory);
+  }
+
+  private <T> void bindStore(Class<T> binding, Class<? extends T> impl) {
+    bind(binding).to(impl);
+    bind(impl).in(Singleton.class);
+    Key<T> key = keyFactory.create(binding);
+    bind(key).to(impl);
+    expose(key);
+  }
+
+  @Override
+  protected void configure() {
+    install(new MyBatisModule() {
+      @Override
+      protected void initialize() {
+        // Ideally, we would install h2 from org.mybatis.guice.datasource.helper.JdbcHelper
+        //     install(JdbcHelper.H2_IN_MEMORY_PRIVATE);
+        // But the in-memory URL is invalid as far as H2 is concerned, so we had to inline
+        // some of the constants here and bind it manually.
+        bindConstant().annotatedWith(named("JDBC.driver")).to(Driver.class.getName());
+        bind(Key.get(String.class, named("JDBC.url"))).toInstance("jdbc:h2:mem:");
+
+        bindDataSourceProviderType(PooledDataSourceProvider.class);
+        bindTransactionFactoryType(JdbcTransactionFactory.class);
+        addMapperClass(LockMapper.class);
+        addMapperClass(JobKeyMapper.class);
+        Properties props = new Properties();
+        // We have no plans to take advantage of multiple DB environments. This is a required
+        // property though, so we use an unnamed environment.
+        props.setProperty("mybatis.environment.id", "");
+        Names.bindProperties(binder(), props);
+        // Full auto-mapping enables population of nested objects with minimal mapper configuration.
+        // Docs on settings can be found here:
+        // http://mybatis.github.io/mybatis-3/configuration.html#settings
+        autoMappingBehavior(AutoMappingBehavior.FULL);
+
+        // TODO(davmclau): ensure that mybatis logging is configured correctly.
+      }
+    });
+    bindStore(LockStore.Mutable.class, DbLockStore.class);
+
+    Key<Storage> storageKey = keyFactory.create(Storage.class);
+    bind(storageKey).to(DbStorage.class);
+    bind(DbStorage.class).in(Singleton.class);
+    expose(storageKey);
+
+    expose(DbStorage.class);
+    expose(SqlSessionFactory.class);
+  }
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/db/DbStorage.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/DbStorage.java b/src/main/java/org/apache/aurora/scheduler/storage/db/DbStorage.java
new file mode 100644
index 0000000..41755c3
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/DbStorage.java
@@ -0,0 +1,162 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db;
+
+import java.io.IOException;
+import java.io.InputStreamReader;
+
+import com.google.common.base.Preconditions;
+import com.google.common.io.CharStreams;
+import com.google.common.util.concurrent.AbstractIdleService;
+import com.google.inject.Inject;
+
+import org.apache.aurora.scheduler.storage.AttributeStore;
+import org.apache.aurora.scheduler.storage.JobStore;
+import org.apache.aurora.scheduler.storage.LockStore;
+import org.apache.aurora.scheduler.storage.QuotaStore;
+import org.apache.aurora.scheduler.storage.SchedulerStore;
+import org.apache.aurora.scheduler.storage.Storage;
+import org.apache.aurora.scheduler.storage.TaskStore;
+import org.apache.ibatis.builder.StaticSqlSource;
+import org.apache.ibatis.exceptions.PersistenceException;
+import org.apache.ibatis.mapping.MappedStatement.Builder;
+import org.apache.ibatis.mapping.SqlCommandType;
+import org.apache.ibatis.session.Configuration;
+import org.apache.ibatis.session.SqlSession;
+import org.apache.ibatis.session.SqlSessionFactory;
+import org.mybatis.guice.transactional.Transactional;
+
+/**
+ * A storage implementation backed by a relational database.
+ *
+ * <p>Delegates read and write concurrency semantics to the underlying database.
+ * In this implementation, {@link #weaklyConsistentRead(Work)} and {@link #consistentRead(Work)}
+ * have identical behaviour as they are both annotated by
+ * {@link org.mybatis.guice.transactional.Transactional}. This class is currently only
+ * partially implemented, with the underlying {@link MutableStoreProvider} only providing
+ * a {@link LockStore.Mutable} implementation. It is designed to be a long term replacement
+ * for {@link org.apache.aurora.scheduler.storage.mem.MemStorage}.</p>
+ */
+class DbStorage extends AbstractIdleService implements Storage {
+
+  private final SqlSessionFactory sessionFactory;
+  private final MutableStoreProvider storeProvider;
+
+  @Inject
+  DbStorage(final SqlSessionFactory sessionFactory, final LockStore.Mutable lockStore) {
+    this.sessionFactory = Preconditions.checkNotNull(sessionFactory);
+    storeProvider = new MutableStoreProvider() {
+      @Override
+      public SchedulerStore.Mutable getSchedulerStore() {
+        throw new UnsupportedOperationException("Not yet implemented.");
+      }
+
+      @Override
+      public JobStore.Mutable getJobStore() {
+        throw new UnsupportedOperationException("Not yet implemented.");
+      }
+
+      @Override
+      public TaskStore getTaskStore() {
+        throw new UnsupportedOperationException("Not yet implemented.");
+      }
+
+      @Override
+      public TaskStore.Mutable getUnsafeTaskStore() {
+        throw new UnsupportedOperationException("Not yet implemented.");
+      }
+
+      @Override
+      public LockStore.Mutable getLockStore() {
+        return lockStore;
+      }
+
+      @Override
+      public QuotaStore.Mutable getQuotaStore() {
+        throw new UnsupportedOperationException("Not yet implemented.");
+      }
+
+      @Override
+      public AttributeStore.Mutable getAttributeStore() {
+        throw new UnsupportedOperationException("Not yet implemented.");
+      }
+    };
+  }
+
+  @Override
+  @Transactional
+  public <T, E extends Exception> T consistentRead(Work<T, E> work) throws StorageException, E {
+    try {
+      return work.apply(storeProvider);
+    } catch (PersistenceException e) {
+      throw new StorageException(e.getMessage(), e);
+    }
+  }
+
+  @Override
+  @Transactional
+  public <T, E extends Exception> T weaklyConsistentRead(Work<T, E> work)
+      throws StorageException, E {
+
+    try {
+      return work.apply(storeProvider);
+    } catch (PersistenceException e) {
+      throw new StorageException(e.getMessage(), e);
+    }
+  }
+
+  @Override
+  @Transactional
+  public <T, E extends Exception> T write(MutateWork<T, E> work) throws StorageException, E {
+    try {
+      return work.apply(storeProvider);
+    } catch (PersistenceException e) {
+      throw new StorageException(e.getMessage(), e);
+    }
+  }
+
+  @Override
+  public void prepare() {
+    startAsync().awaitRunning();
+  }
+
+  /**
+   * Creates the SQL schema during service start-up.
+   * Note: This design assumes a volatile database engine.
+   */
+  @Override
+  @Transactional
+  protected void startUp() throws IOException {
+    Configuration configuration = sessionFactory.getConfiguration();
+    String createStatementName = "create_tables";
+    configuration.addMappedStatement(new Builder(
+        configuration,
+        createStatementName,
+        new StaticSqlSource(
+            configuration,
+            CharStreams.toString(
+                new InputStreamReader(DbStorage.class.getResourceAsStream("schema.sql"), "UTF-8"))),
+        SqlCommandType.UPDATE)
+        .build());
+
+    try (SqlSession session = sessionFactory.openSession()) {
+      session.update(createStatementName);
+    }
+  }
+
+  @Override
+  protected void shutDown() {
+    // noop
+  }
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/db/JobKeyMapper.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/JobKeyMapper.java b/src/main/java/org/apache/aurora/scheduler/storage/db/JobKeyMapper.java
new file mode 100644
index 0000000..844714c
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/JobKeyMapper.java
@@ -0,0 +1,40 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db;
+
+import java.util.List;
+
+import org.apache.aurora.gen.JobKey;
+
+/**
+ * MyBatis mapper class for JobKeyMapper.xml
+ *
+ * See http://mybatis.github.io/mybatis-3/sqlmap-xml.html for more details.
+ */
+interface JobKeyMapper {
+  /**
+   * Inserts a new job key into the database.
+   */
+  void insert(JobKey key);
+
+  /**
+   * Permanently deletes a job key from the database.
+   */
+  void delete(JobKey key);
+
+  /**
+   * Selects all job keys from the database.
+   */
+  List<JobKey> selectAll();
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/db/LockKeyMapper.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/LockKeyMapper.java b/src/main/java/org/apache/aurora/scheduler/storage/db/LockKeyMapper.java
new file mode 100644
index 0000000..72543bf
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/LockKeyMapper.java
@@ -0,0 +1,49 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db;
+
+import com.google.inject.Inject;
+
+import org.apache.aurora.gen.LockKey;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+/**
+ * Mapper for LockKeys. Not a MyBatis mapper, this just encapsulates the logic for writing
+ * union types so it does not leak into the related object's implementation.
+ *
+ * TODO(davmclau):
+ * Consider creating these classes with code generation since something like this will be needed
+ * for all union field relationships. Might not be possible unless the code generator also defines
+ * a mapper interface for every field in the union as well as the associated XML mapper config
+ * with the SQL to satisfy the interface.
+ *
+ */
+public class LockKeyMapper {
+
+  private final JobKeyMapper jobKeyMapper;
+
+  @Inject
+  public LockKeyMapper(JobKeyMapper jobKeyMapper) {
+    this.jobKeyMapper = checkNotNull(jobKeyMapper);
+  }
+
+  public void insert(LockKey key) {
+    if (key.isSetJob()) {
+      jobKeyMapper.insert(checkNotNull(key.getJob()));
+    } else {
+      throw new IllegalArgumentException("Unsupported lock type on LockKey.");
+    }
+  }
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/db/LockMapper.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/LockMapper.java b/src/main/java/org/apache/aurora/scheduler/storage/db/LockMapper.java
new file mode 100644
index 0000000..b9f338c
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/LockMapper.java
@@ -0,0 +1,53 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db;
+
+import java.util.List;
+
+import org.apache.aurora.gen.Lock;
+import org.apache.aurora.gen.LockKey;
+import org.apache.aurora.scheduler.storage.db.views.LockRow;
+
+/**
+ * MyBatis mapper class for LockMapper.xml
+ *
+ * See http://mybatis.github.io/mybatis-3/sqlmap-xml.html for more details.
+ */
+interface LockMapper {
+
+  /**
+   * Inserts a lock into the database.
+   */
+  void insert(Lock lock);
+
+  /**
+   * Deletes all locks from the database with the given lockKey.
+   */
+  void delete(LockKey lockKey);
+
+  /**
+   * Deletes all locks from the database.
+   */
+  void truncate();
+
+  /**
+   * Selects all locks from the database as {@link LockRow} instances.
+   */
+  List<LockRow> selectAll();
+
+  /**
+   * Fetches the lock with the given lock key.
+   */
+  LockRow select(LockKey lockKey);
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/db/MigrationModule.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/MigrationModule.java b/src/main/java/org/apache/aurora/scheduler/storage/db/MigrationModule.java
new file mode 100644
index 0000000..7e98ebf
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/MigrationModule.java
@@ -0,0 +1,48 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db;
+
+import com.google.inject.AbstractModule;
+import com.twitter.common.inject.Bindings.KeyFactory;
+
+import org.apache.aurora.scheduler.storage.LockStore;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+/**
+ * Temporary module to wire the two partial storage implementations together as we
+ * migrate from MemStorage to DbStorage.  This accepts two {@link KeyFactory}s,
+ * one that references the binding scope for the feature-complete write-behind
+ * volatile storage system, and one for the binding scope of the new and partially-implemented
+ * storage system.
+ * <p>
+ *  Once the new storage system is feature-complete, this module will be deleted
+ *  as the binding bridge is no longer necessary.
+ * </p>
+ */
+public class MigrationModule extends AbstractModule {
+
+  private final KeyFactory toFactory;
+  private final KeyFactory fromFactory;
+
+  public MigrationModule(KeyFactory from, KeyFactory to) {
+    this.fromFactory = checkNotNull(from);
+    this.toFactory = checkNotNull(to);
+  }
+
+  @Override
+  protected void configure() {
+    bind(fromFactory.create(LockStore.Mutable.class)).to(toFactory.create(LockStore.Mutable.class));
+  }
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/db/views/LockRow.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/db/views/LockRow.java b/src/main/java/org/apache/aurora/scheduler/storage/db/views/LockRow.java
new file mode 100644
index 0000000..aaa0a68
--- /dev/null
+++ b/src/main/java/org/apache/aurora/scheduler/storage/db/views/LockRow.java
@@ -0,0 +1,46 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db.views;
+
+import org.apache.aurora.gen.JobKey;
+import org.apache.aurora.gen.Lock;
+import org.apache.aurora.gen.LockKey;
+
+/**
+ * The union of {@link Lock} and {@link JobKey}. This class is required because the generated
+ * thrift code for {@code union} fields is incompatible with the mapping behaviour of MyBatis.
+ * This class works around this incompatibility by explicitly initialising an empty lock with
+ * the union field set to an empty {@link JobKey} instance and exposing a getter method for MyBatis.
+ * Note that this only works because the {@link LockKey} is a union of exactly one type. If LockKey
+ * is modified to support more types, we will need to rework this design.
+ *
+ * TODO(davmclau):
+ * These intermediate classes for resolving relationships on Thrift union types should be
+ * auto-generated, as the logic will be identical in each one. The generated code needs to allow
+ * for exactly one field in the union to be set, returning null when the getter for the field
+ * is called the first time.
+ */
+public class LockRow {
+  private final Lock lock = new Lock();
+
+  public LockRow() {
+    LockKey lockKey = new LockKey();
+    lock.setKey(lockKey);
+    lockKey.setJob(new JobKey());
+  }
+
+  public Lock getLock() {
+    return lock;
+  }
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java b/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java
index 7fcc072..39f4712 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorage.java
@@ -292,6 +292,7 @@ public class LogStorage implements NonVolatileStorage, DistributedSnapshotStore
 
   @Override
   public synchronized void prepare() {
+    writeBehindStorage.prepare();
     // Open the log to make a log replica available to the scheduler group.
     try {
       streamManager = logManager.open();

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorageModule.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorageModule.java b/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorageModule.java
index 96da989..23ee32b 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorageModule.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/log/LogStorageModule.java
@@ -85,8 +85,8 @@ public class LogStorageModule extends AbstractModule {
         .toInstance(MAX_LOG_ENTRY_SIZE.get());
     bind(LogManager.class).in(Singleton.class);
     bind(Boolean.class).annotatedWith(SnapshotSetting.class).toInstance(DEFLATE_SNAPSHOTS.get());
-
     bind(LogStorage.class).in(Singleton.class);
+
     install(CallOrderEnforcingStorage.wrappingModule(LogStorage.class));
     bind(DistributedSnapshotStore.class).to(LogStorage.class);
   }

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/mem/MemLockStore.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/mem/MemLockStore.java b/src/main/java/org/apache/aurora/scheduler/storage/mem/MemLockStore.java
deleted file mode 100644
index 1daab63..0000000
--- a/src/main/java/org/apache/aurora/scheduler/storage/mem/MemLockStore.java
+++ /dev/null
@@ -1,58 +0,0 @@
-/**
- * Licensed 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.aurora.scheduler.storage.mem;
-
-import java.util.Map;
-import java.util.Set;
-
-import com.google.common.base.Optional;
-import com.google.common.collect.ImmutableSet;
-import com.google.common.collect.Maps;
-
-import org.apache.aurora.scheduler.storage.LockStore;
-import org.apache.aurora.scheduler.storage.entities.ILock;
-import org.apache.aurora.scheduler.storage.entities.ILockKey;
-
-/**
- * An in-memory lock store.
- */
-class MemLockStore implements LockStore.Mutable {
-
-  private final Map<ILockKey, ILock> locks = Maps.newConcurrentMap();
-
-  @Override
-  public void saveLock(ILock lock) {
-    locks.put(lock.getKey(), lock);
-  }
-
-  @Override
-  public void removeLock(ILockKey lockKey) {
-    locks.remove(lockKey);
-  }
-
-  @Override
-  public void deleteLocks() {
-    locks.clear();
-  }
-
-  @Override
-  public Set<ILock> fetchLocks() {
-    return ImmutableSet.copyOf(locks.values());
-  }
-
-  @Override
-  public Optional<ILock> fetchLock(ILockKey lockKey) {
-    return Optional.fromNullable(locks.get(lockKey));
-  }
-}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorage.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorage.java b/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorage.java
index 4a7db20..90d9a65 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorage.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorage.java
@@ -13,11 +13,20 @@
  */
 package org.apache.aurora.scheduler.storage.mem;
 
+import java.lang.annotation.ElementType;
+import java.lang.annotation.Retention;
+import java.lang.annotation.RetentionPolicy;
+import java.lang.annotation.Target;
 import java.util.concurrent.atomic.AtomicLong;
 
 import javax.inject.Inject;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.inject.BindingAnnotation;
+import com.google.inject.Guice;
+import com.google.inject.Injector;
+import com.google.inject.Key;
+import com.twitter.common.inject.Bindings;
 import com.twitter.common.inject.TimedInterceptor.Timed;
 import com.twitter.common.stats.SlidingStats;
 import com.twitter.common.stats.StatImpl;
@@ -32,6 +41,7 @@ import org.apache.aurora.scheduler.storage.ReadWriteLockManager.LockType;
 import org.apache.aurora.scheduler.storage.SchedulerStore;
 import org.apache.aurora.scheduler.storage.Storage;
 import org.apache.aurora.scheduler.storage.TaskStore;
+import org.apache.aurora.scheduler.storage.db.DbModule;
 
 import static com.google.common.base.Preconditions.checkNotNull;
 
@@ -58,16 +68,26 @@ public class MemStorage implements Storage {
 
   private final MutableStoreProvider storeProvider;
   private final ReadWriteLockManager lockManager = new ReadWriteLockManager();
+  private final Storage delegatedStore;
+
+  /**
+   * Identifies a storage layer to be delegated to instead of mem storage.
+   */
+  @Retention(RetentionPolicy.RUNTIME)
+  @Target({ ElementType.PARAMETER, ElementType.METHOD })
+  @BindingAnnotation
+  public @interface Delegated { }
 
   @Inject
   MemStorage(
       final SchedulerStore.Mutable schedulerStore,
       final JobStore.Mutable jobStore,
       final TaskStore.Mutable taskStore,
-      final LockStore.Mutable lockStore,
+      @Delegated final LockStore.Mutable lockStore,
+      @Delegated final Storage delegated,
       final QuotaStore.Mutable quotaStore,
       final AttributeStore.Mutable attributeStore) {
-
+    this.delegatedStore = delegated;
     storeProvider = new MutableStoreProvider() {
       @Override
       public SchedulerStore.Mutable getSchedulerStore() {
@@ -117,14 +137,14 @@ public class MemStorage implements Storage {
    * Creates a new empty in-memory storage for use in testing.
    */
   @VisibleForTesting
-  public static MemStorage newEmptyStorage() {
-    return new MemStorage(
-        new MemSchedulerStore(),
-        new MemJobStore(),
-        new MemTaskStore(),
-        new MemLockStore(),
-        new MemQuotaStore(),
-        new MemAttributeStore());
+  public static Storage newEmptyStorage() {
+    Injector injector = Guice.createInjector(
+        new DbModule(Bindings.annotatedKeyFactory(Delegated.class)),
+        new MemStorageModule(Bindings.annotatedKeyFactory(Volatile.class)));
+
+    Storage storage = injector.getInstance(Key.get(Storage.class, Volatile.class));
+    storage.prepare();
+    return storage;
   }
 
   private <S extends StoreProvider, T, E extends Exception> T doWork(
@@ -163,6 +183,11 @@ public class MemStorage implements Storage {
     return doWork(LockType.WRITE, storeProvider, work, writeStats, writeLockWaitNanos);
   }
 
+  @Override
+  public void prepare() throws StorageException {
+    delegatedStore.prepare();
+  }
+
   @Timed("mem_storage_weakly_consistent_read_operation")
   @Override
   public <T, E extends Exception> T weaklyConsistentRead(Work<T, E> work)

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorageModule.java
----------------------------------------------------------------------
diff --git a/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorageModule.java b/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorageModule.java
index 6137f5b..acafd40 100644
--- a/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorageModule.java
+++ b/src/main/java/org/apache/aurora/scheduler/storage/mem/MemStorageModule.java
@@ -21,7 +21,6 @@ import com.twitter.common.inject.Bindings.KeyFactory;
 
 import org.apache.aurora.scheduler.storage.AttributeStore;
 import org.apache.aurora.scheduler.storage.JobStore;
-import org.apache.aurora.scheduler.storage.LockStore;
 import org.apache.aurora.scheduler.storage.QuotaStore;
 import org.apache.aurora.scheduler.storage.SchedulerStore;
 import org.apache.aurora.scheduler.storage.Storage;
@@ -76,7 +75,6 @@ public final class MemStorageModule extends PrivateModule {
     bindStore(SchedulerStore.Mutable.class, MemSchedulerStore.class);
     bindStore(JobStore.Mutable.class, MemJobStore.class);
     bindStore(TaskStore.Mutable.class, MemTaskStore.class);
-    bindStore(LockStore.Mutable.class, MemLockStore.class);
     bindStore(QuotaStore.Mutable.class, MemQuotaStore.class);
     bindStore(AttributeStore.Mutable.class, MemAttributeStore.class);
   }

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/resources/org/apache/aurora/scheduler/storage/db/JobKeyMapper.xml
----------------------------------------------------------------------
diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/JobKeyMapper.xml b/src/main/resources/org/apache/aurora/scheduler/storage/db/JobKeyMapper.xml
new file mode 100644
index 0000000..9a87d77
--- /dev/null
+++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/JobKeyMapper.xml
@@ -0,0 +1,33 @@
+<?xml version="1.0" encoding="UTF-8" ?>
+<!DOCTYPE mapper
+    PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN"
+    "http://mybatis.org/dtd/mybatis-3-mapper.dtd">
+<mapper namespace="org.apache.aurora.scheduler.storage.db.JobKeyMapper">
+  <insert id="insert">
+    INSERT INTO job_keys (
+      role,
+      environment,
+      name
+    ) VALUES (
+      #{role},
+      #{environment},
+      #{name}
+    )
+  </insert>
+  <sql id="job_key_clause">
+    role = #{role}
+    AND environment = #{environment}
+    AND name = #{name}
+  </sql>
+  <delete id="delete">
+    DELETE FROM job_keys
+    WHERE <include refid="job_key_clause"/>
+  </delete>
+  <select id="exists" resultType="org.apache.aurora.gen.JobKey">
+    SELECT * FROM job_keys
+    WHERE <include refid="job_key_clause"/>
+  </select>
+  <select id="selectAll" resultType="org.apache.aurora.gen.JobKey">
+    SELECT * FROM job_keys
+  </select>
+</mapper>

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/resources/org/apache/aurora/scheduler/storage/db/LockMapper.xml
----------------------------------------------------------------------
diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/LockMapper.xml b/src/main/resources/org/apache/aurora/scheduler/storage/db/LockMapper.xml
new file mode 100644
index 0000000..3797e4c
--- /dev/null
+++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/LockMapper.xml
@@ -0,0 +1,69 @@
+<?xml version="1.0" encoding="UTF-8" ?>
+<!DOCTYPE mapper
+    PUBLIC "-//mybatis.org//DTD Mapper 3.0//EN"
+    "http://mybatis.org/dtd/mybatis-3-mapper.dtd">
+<mapper namespace="org.apache.aurora.scheduler.storage.db.LockMapper">
+  <insert id="insert">
+    INSERT INTO locks (
+      job_key_id,
+      token,
+      user,
+      timestampMs,
+      message
+    ) VALUES (
+      (
+        SELECT ID
+        FROM job_keys
+        WHERE role = #{key.job.role}
+          AND environment = #{key.job.environment}
+          AND name = #{key.job.name}
+      ),
+      #{token},
+      #{user},
+      #{timestampMs},
+      #{message}
+    )
+  </insert>
+
+  <resultMap id="lockResultMap" type="org.apache.aurora.scheduler.storage.db.views.LockRow">
+    <!--
+      Normally you can have MyBatis auto-map these columns and/or use assocations, but
+      in this case we need to work-around thrift union type issues and prevent MyBatis
+      from trying to create new objects like it does with associations.
+      Explicitly mapping each column to a single base property (lock, here) does just that.
+    -->
+    <result column="token" property="lock.token"/>
+    <result column="user" property="lock.user"/>
+    <result column="timestampMs" property="lock.timestampMs"/>
+    <result column="message" property="lock.message"/>
+    <result column="role" property="lock.key.job.role"/>
+    <result column="environment" property="lock.key.job.environment"/>
+    <result column="name" property="lock.key.job.name"/>
+  </resultMap>
+
+  <select id="selectAll" resultMap="lockResultMap">
+    SELECT * FROM locks AS lock
+    JOIN job_keys AS key ON job_key_id = key.id
+  </select>
+  <sql id="jobKeyScope">
+    OUTER JOIN job_keys AS key ON key.role = #{job.role}
+    AND key.environment = #{job.environment}
+    AND key.name = #{job.name}
+    AND key.id = job_key_id
+  </sql>
+  <select id="select" resultMap="lockResultMap">
+    SELECT * FROM locks <include refid="jobKeyScope" />
+  </select>
+  <delete id="delete">
+    DELETE FROM locks
+    WHERE job_key_id
+      IN (SELECT id
+          FROM job_keys
+          WHERE role = #{job.role}
+            AND environment = #{job.environment}
+            AND name = #{job.name})
+  </delete>
+  <delete id="truncate">
+    DELETE FROM locks
+  </delete>
+</mapper>

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql
----------------------------------------------------------------------
diff --git a/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql b/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql
new file mode 100644
index 0000000..405fda5
--- /dev/null
+++ b/src/main/resources/org/apache/aurora/scheduler/storage/db/schema.sql
@@ -0,0 +1,21 @@
+-- schema for h2 engine.
+
+CREATE TABLE job_keys(
+  id INT IDENTITY,
+  role VARCHAR NOT NULL,
+  environment VARCHAR NOT NULL,
+  name VARCHAR NOT NULL,
+
+  UNIQUE(role, environment, name)
+);
+
+CREATE TABLE locks(
+  id INT IDENTITY,
+  job_key_id INT NOT NULL REFERENCES job_keys(id),
+  token VARCHAR NOT NULL,
+  user VARCHAR NOT NULL,
+  timestampMs BIGINT NOT NULL,
+  message VARCHAR,
+
+  UNIQUE(job_key_id)
+);

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/test/java/org/apache/aurora/scheduler/storage/db/DbLockStoreTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/storage/db/DbLockStoreTest.java b/src/test/java/org/apache/aurora/scheduler/storage/db/DbLockStoreTest.java
new file mode 100644
index 0000000..3a50454
--- /dev/null
+++ b/src/test/java/org/apache/aurora/scheduler/storage/db/DbLockStoreTest.java
@@ -0,0 +1,203 @@
+/**
+ * Licensed 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.aurora.scheduler.storage.db;
+
+import java.io.IOException;
+import java.util.Set;
+
+import com.google.common.base.Optional;
+import com.google.common.collect.ImmutableSet;
+import com.google.inject.Guice;
+import com.google.inject.Injector;
+import com.twitter.common.inject.Bindings;
+
+import org.apache.aurora.gen.JobKey;
+import org.apache.aurora.gen.Lock;
+import org.apache.aurora.gen.LockKey;
+import org.apache.aurora.scheduler.base.JobKeys;
+import org.apache.aurora.scheduler.storage.Storage;
+import org.apache.aurora.scheduler.storage.Storage.MutableStoreProvider;
+import org.apache.aurora.scheduler.storage.Storage.MutateWork;
+import org.apache.aurora.scheduler.storage.Storage.StorageException;
+import org.apache.aurora.scheduler.storage.Storage.StoreProvider;
+import org.apache.aurora.scheduler.storage.Storage.Work.Quiet;
+import org.apache.aurora.scheduler.storage.entities.ILock;
+import org.apache.aurora.scheduler.storage.entities.ILockKey;
+import org.junit.Before;
+import org.junit.Test;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.fail;
+
+public class DbLockStoreTest {
+
+  private DbStorage storage;
+
+  private void assertLocks(final ILock... expected) {
+    assertEquals(
+        ImmutableSet.<ILock>builder().add(expected).build(),
+        storage.consistentRead(new Quiet<Set<ILock>>() {
+          @Override
+          public Set<ILock> apply(Storage.StoreProvider storeProvider) {
+            return storeProvider.getLockStore().fetchLocks();
+          }
+        }));
+  }
+
+  private Optional<ILock> getLock(final ILockKey key) {
+    return storage.consistentRead(new Quiet<Optional<ILock>>() {
+      @Override
+      public Optional<ILock> apply(StoreProvider storeProvider) {
+        return storeProvider.getLockStore().fetchLock(key);
+      }
+    });
+  }
+
+  private void saveLocks(final ILock... locks) {
+    storage.write(new MutateWork.Quiet<Void>() {
+      @Override
+      public Void apply(MutableStoreProvider storeProvider) {
+        for (ILock lock : locks) {
+          storeProvider.getLockStore().saveLock(lock);
+        }
+        return null;
+      }
+    });
+  }
+
+  private void removeLocks(final ILock... locks) {
+    storage.write(new MutateWork.Quiet<Void>() {
+      @Override
+      public Void apply(MutableStoreProvider storeProvider) {
+        for (ILock lock : locks) {
+          storeProvider.getLockStore().removeLock(lock.getKey());
+        }
+        return null;
+      }
+    });
+  }
+
+  private static ILock makeLock(JobKey key) {
+    return ILock.build(new Lock()
+      .setKey(LockKey.job(key))
+      .setToken("lock1")
+      .setUser("testUser")
+      .setMessage("Test message")
+      .setTimestampMs(12345L));
+  }
+
+  @Before
+  public void setUp() throws IOException {
+    Injector injector = Guice.createInjector(new DbModule(Bindings.KeyFactory.PLAIN));
+    storage = injector.getInstance(DbStorage.class);
+    storage.prepare();
+  }
+
+  @Test
+  public void testLocks() throws Exception {
+    assertLocks();
+
+    String role = "testRole";
+    String env = "testEnv";
+    String job1 = "testJob1";
+    String job2 = "testJob2";
+
+    ILock lock1 = makeLock(JobKeys.from(role, env, job1).newBuilder());
+    ILock lock2 = makeLock(JobKeys.from(role, env, job2).newBuilder());
+
+    saveLocks(lock1, lock2);
+    assertLocks(lock1, lock2);
+    removeLocks(lock1);
+
+    assertLocks(lock2);
+  }
+
+  @Test
+  public void testRepeatedWrite() throws Exception {
+    assertLocks();
+
+    String role = "testRole";
+    String env = "testEnv";
+    String job = "testJob";
+
+    ILock lock = makeLock(JobKeys.from(role, env, job).newBuilder());
+
+    saveLocks(lock);
+    try {
+      saveLocks(lock);
+      fail("saveLock should have failed unique constraint check.");
+    } catch (StorageException e) {
+      // expected
+    }
+
+    assertLocks(lock);
+  }
+
+  @Test
+  public void testExistingJobKey() throws Exception {
+    String role = "testRole";
+    String env = "testEnv";
+    String job = "testJob";
+
+    ILock lock = makeLock(JobKeys.from(role, env, job).newBuilder());
+
+    saveLocks(lock);
+    removeLocks(lock);
+    saveLocks(lock);
+
+    assertLocks(lock);
+  }
+
+  @Test
+  public void testGetLock() throws Exception {
+    assertLocks();
+
+    String role = "testRole";
+    String env = "testEnv";
+    String job = "testJob";
+
+    final ILock lock = makeLock(JobKeys.from(role, env, job).newBuilder());
+
+    assertEquals(Optional.absent(), getLock(lock.getKey()));
+
+    saveLocks(lock);
+    assertEquals(Optional.<ILock>of(lock), getLock(lock.getKey()));
+  }
+
+  @Test
+  public void testDeleteAllLocks() throws Exception {
+    assertLocks();
+
+    String role = "testRole";
+    String env = "testEnv";
+    String job1 = "testJob1";
+    String job2 = "testJob2";
+
+    ILock lock1 = makeLock(JobKeys.from(role, env, job1).newBuilder());
+    ILock lock2 = makeLock(JobKeys.from(role, env, job2).newBuilder());
+
+    saveLocks(lock1, lock2);
+    assertLocks(lock1, lock2);
+
+    storage.write(new MutateWork.Quiet<Void>() {
+      @Override
+      public Void apply(MutableStoreProvider storeProvider) {
+        storeProvider.getLockStore().deleteLocks();
+        return null;
+      }
+    });
+
+    assertLocks();
+  }
+}

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java b/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java
index 6f6d9a4..53e5749 100644
--- a/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java
+++ b/src/test/java/org/apache/aurora/scheduler/storage/log/LogStorageTest.java
@@ -134,6 +134,8 @@ public class LogStorageTest extends EasyMockTest {
     stream = createMock(Stream.class);
     streamMatcher = LogOpMatcher.matcherFor(stream);
     position = createMock(Position.class);
+
+    storageUtil.storage.prepare();
   }
 
   @Test

http://git-wip-us.apache.org/repos/asf/incubator-aurora/blob/1b6a55e4/src/test/java/org/apache/aurora/scheduler/storage/mem/MemLockStoreTest.java
----------------------------------------------------------------------
diff --git a/src/test/java/org/apache/aurora/scheduler/storage/mem/MemLockStoreTest.java b/src/test/java/org/apache/aurora/scheduler/storage/mem/MemLockStoreTest.java
deleted file mode 100644
index d4fe607..0000000
--- a/src/test/java/org/apache/aurora/scheduler/storage/mem/MemLockStoreTest.java
+++ /dev/null
@@ -1,77 +0,0 @@
-/**
- * Licensed 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.aurora.scheduler.storage.mem;
-
-import com.google.common.base.Optional;
-import com.google.common.collect.ImmutableSet;
-
-import org.apache.aurora.gen.Lock;
-import org.apache.aurora.gen.LockKey;
-import org.apache.aurora.scheduler.base.JobKeys;
-import org.apache.aurora.scheduler.storage.LockStore;
-import org.apache.aurora.scheduler.storage.entities.ILock;
-import org.apache.aurora.scheduler.storage.entities.ILockKey;
-import org.junit.Before;
-import org.junit.Test;
-
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertNotNull;
-
-public class MemLockStoreTest {
-  private LockStore.Mutable store;
-
-  @Before
-  public void setUp() {
-    store = new MemLockStore();
-  }
-
-  @Test
-  public void testLocks() {
-    final String role = "testRole";
-    final String env = "testEnv";
-    final String job1 = "testJob1";
-    final String job2 = "testJob2";
-    ILock lock1 = ILock.build(new Lock()
-        .setKey(LockKey.job(JobKeys.from(role, env, job1).newBuilder()))
-        .setToken("lock1")
-        .setUser("testUser")
-        .setTimestampMs(12345L));
-    ILock lock2 = ILock.build(new Lock()
-        .setKey(LockKey.job(JobKeys.from(role, env, job2).newBuilder()))
-        .setToken("lock2")
-        .setUser("testUser")
-        .setTimestampMs(12345L)
-        .setMessage("Test message"));
-
-    store.saveLock(lock1);
-    store.saveLock(lock2);
-
-    assertEquals(Optional.of(lock1),
-        store.fetchLock(ILockKey.build(LockKey.job(JobKeys.from(role, env, job1).newBuilder()))));
-    assertEquals(Optional.of(lock2),
-        store.fetchLock(ILockKey.build(LockKey.job(JobKeys.from(role, env, job2).newBuilder()))));
-    assertEquals(ImmutableSet.of(lock1, lock2), store.fetchLocks());
-
-    store.removeLock(ILockKey.build(LockKey.job(JobKeys.from(role, env, job1).newBuilder())));
-    assertEquals(Optional.<ILock>absent(),
-        store.fetchLock(ILockKey.build(LockKey.job(JobKeys.from(role, env, job1).newBuilder()))));
-
-    assertEquals(Optional.of(lock2),
-        store.fetchLock(ILockKey.build(LockKey.job(JobKeys.from(role, env, job2).newBuilder()))));
-    assertNotNull(
-        store.fetchLock(ILockKey.build(LockKey.job(JobKeys.from(role, env, job2).newBuilder())))
-        .get().getMessage());
-    assertEquals(ImmutableSet.of(lock2), store.fetchLocks());
-  }
-}


Mime
View raw message