Initial (simple) lock implementation

just a brain dump of a distributed cassandra lock from (http://www.datastax.com/dev/blog/consensus-on-cassandra) will get some cleanup the other day

Author: larsp

lack/src/main/java/io/datanerds/lack/Lack.java

@@ -0,0 +1,50 @@
+package io.datanerds.lack;
+
+import com.datastax.driver.core.ResultSet;
+import com.datastax.driver.core.Session;
+import io.datanerds.lack.cassandra.CassandraClient;
+import io.datanerds.lack.cassandra.LackConfig;
+import io.datanerds.lack.cassandra.Statements;
+
+public class Lack {
+
+    private final String owner;
+    private final CassandraClient client;
+    private final Session session;
+    private final Statements statements;
+
+    public Lack(LackConfig lackConfig, String owner) {
+        this.owner = owner;
+        client = new CassandraClient();
+        session = client.init(lackConfig);
+        statements = new Statements(session);
+        statements.create();
+        statements.setTtl(lackConfig.ttlInSeconds);
+    }
+
+    public void acquire(String resource) {
+        ResultSet result = statements.acquire(resource, owner);
+        if (!result.wasApplied()) {
+            throw new LackException(String.format("Could not acquire lock for '%s'", resource));
+        }
+    }
+
+    public void renew(String resource) {
+        ResultSet result = statements.renew(resource, owner);
+        if (!result.wasApplied()) {
+            throw new LackException(String.format("Could not renew lock for '%s'", resource));
+        }
+    }
+
+    public void release(String resource) {
+        ResultSet result = statements.release(resource, owner);
+        if (!result.wasApplied()) {
+            throw new LackException(String.format("Could not release lock for '%s'", resource));
+        }
+    }
+
+    public void stop() {
+        session.close();
+        client.close();
+    }
+}

lack/src/main/java/io/datanerds/lack/LackException.java

@@ -0,0 +1,7 @@
+package io.datanerds.lack;
+
+public class LackException extends RuntimeException {
+    public LackException(String message) {
+        super(message);
+    }
+}

lack/src/main/java/io/datanerds/lack/cassandra/CassandraClient.java

@@ -0,0 +1,41 @@
+package io.datanerds.lack.cassandra;
+
+import com.datastax.driver.core.Cluster;
+import com.datastax.driver.core.Metadata;
+import com.datastax.driver.core.Session;
+import com.datastax.driver.core.SocketOptions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class CassandraClient {
+
+    private static final Logger logger = LoggerFactory.getLogger(CassandraClient.class);
+    private Cluster cluster;
+
+    public Session init(LackConfig config) {
+        Cluster.Builder clusterBuilder = new Cluster.Builder().addContactPoints(config.nodes)
+                .withSocketOptions(new SocketOptions().setTcpNoDelay(true).setReuseAddress(true));
+        if (config.username != null && config.password != null) {
+            clusterBuilder.withCredentials(config.username, config.password);
+        }
+
+        this.cluster = clusterBuilder.build();
+        logClusterInfo();
+        Session session = cluster.connect(config.keyspace);
+        return session;
+    }
+
+    public void close() {
+        cluster.close();
+    }
+
+    private void logClusterInfo() {
+        Metadata metadata = cluster.getMetadata();
+        logger.info("Connected to cluster: {}; Protocol Version: {}", metadata.getClusterName(),
+                cluster.getConfiguration().getProtocolOptions().getProtocolVersionEnum());
+
+        metadata.getAllHosts()
+                .forEach(host -> logger.info("Added '{}' from DC '{}' which is running '{}'", host.getAddress(),
+                        host.getDatacenter(), host.getCassandraVersion().toString()));
+    }
+}

lack/src/main/java/io/datanerds/lack/cassandra/Constants.java

@@ -0,0 +1,10 @@
+package io.datanerds.lack.cassandra;
+
+public interface Constants {
+    String TABLE = "leases";
+
+    interface Columns {
+        String RESOURCE = "resource";
+        String OWNER = "owner";
+    }
+}

lack/src/main/java/io/datanerds/lack/cassandra/LackConfig.java

@@ -0,0 +1,18 @@
+package io.datanerds.lack.cassandra;
+
+public class LackConfig {
+
+    public final String username;
+    public final String password;
+    public final String[] nodes;
+    public final String keyspace;
+    public final int ttlInSeconds;
+
+    public LackConfig(String username, String password, String[] nodes, String keyspace, int ttlInSeconds) {
+        this.username = username;
+        this.password = password;
+        this.nodes = nodes;
+        this.keyspace = keyspace;
+        this.ttlInSeconds = ttlInSeconds;
+    }
+}

lack/src/main/java/io/datanerds/lack/cassandra/Statements.java

@@ -0,0 +1,77 @@
+package io.datanerds.lack.cassandra;
+
+import com.datastax.driver.core.BoundStatement;
+import com.datastax.driver.core.PreparedStatement;
+import com.datastax.driver.core.ResultSet;
+import com.datastax.driver.core.Session;
+import com.datastax.driver.core.querybuilder.BuiltStatement;
+import com.datastax.driver.core.querybuilder.Delete;
+import com.datastax.driver.core.querybuilder.Insert;
+import com.datastax.driver.core.querybuilder.Update;
+import com.datastax.driver.core.schemabuilder.Create;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import static com.datastax.driver.core.DataType.text;
+import static com.datastax.driver.core.querybuilder.QueryBuilder.*;
+import static com.datastax.driver.core.schemabuilder.SchemaBuilder.alterTable;
+import static com.datastax.driver.core.schemabuilder.SchemaBuilder.createTable;
+import static io.datanerds.lack.cassandra.Constants.Columns.OWNER;
+import static io.datanerds.lack.cassandra.Constants.Columns.RESOURCE;
+import static io.datanerds.lack.cassandra.Constants.TABLE;
+
+public class Statements {
+
+    private static final Logger logger = LoggerFactory.getLogger(Statements.class);
+    private final Map<BuiltStatement, PreparedStatement> statements = new ConcurrentHashMap<>();
+    private final Session session;
+
+    private final Create create = createTable(TABLE).addPartitionKey(RESOURCE, text())
+            .addColumn(OWNER, text())
+            .ifNotExists();
+
+    private final Insert acquire = insertInto(TABLE).value(RESOURCE, bindMarker())
+            .value(OWNER, bindMarker())
+            .ifNotExists();
+
+    private final Update.Conditions renew = update(TABLE).with(set(OWNER, bindMarker()))
+            .where(eq(RESOURCE, bindMarker()))
+            .onlyIf(eq(OWNER, bindMarker()));
+
+    private final Delete.Conditions release = delete().from(TABLE)
+            .where(eq(RESOURCE, bindMarker()))
+            .onlyIf(eq(OWNER, bindMarker()));
+
+    public Statements(Session session) {
+        this.session = session;
+    }
+
+    public void create() {
+        session.execute(create);
+    }
+
+    public void setTtl(int ttl) {
+        logger.info("Setting a default TTL of '{}' seconds", ttl);
+        session.execute(alterTable("leases").withOptions().defaultTimeToLive(ttl));
+    }
+
+    public ResultSet acquire(String resource, String owner) {
+        BoundStatement boundStatement = statements.computeIfAbsent(acquire, session::prepare).bind(resource, owner);
+        return session.execute(boundStatement);
+    }
+
+    public ResultSet renew(String resource, String owner) {
+        BoundStatement boundStatement = statements.computeIfAbsent(renew, session::prepare)
+                .bind(owner, resource, owner);
+        return session.execute(boundStatement);
+    }
+
+    public ResultSet release(String resource, String owner) {
+        BoundStatement boundStatement = statements.computeIfAbsent(release, session::prepare).bind(resource, owner);
+        return session.execute(boundStatement);
+
+    }
+}

lack/src/test/java/io/datanerds/lack/LackTest.java

@@ -0,0 +1,70 @@
+package io.datanerds.lack;
+
+import io.datanerds.lack.cassandra.LackConfig;
+import org.junit.*;
+import org.junit.rules.ExpectedException;
+
+import java.util.UUID;
+
+public class LackTest {
+
+    private static LackConfig config = new LackConfig(null, null, new String[]{"127.0.0.1"}, "lack", 1);
+    private static Lack lack;
+    private static Lack otherLack;
+    private String resource;
+
+    @Rule
+    public ExpectedException thrown = ExpectedException.none();
+
+    @BeforeClass
+    public static void setup() {
+        lack = new Lack(config, "lack");
+        otherLack = new Lack(config, "otherLack");
+    }
+
+    @AfterClass
+    public static void tearDown() {
+        lack.stop();
+    }
+
+    @Before
+    public void setupResource() {
+        this.resource = UUID.randomUUID().toString();
+    }
+
+    @Test
+    public void simpleAcquireRenewAndRelease() {
+        lack.acquire(resource);
+        lack.renew(resource);
+        lack.release(resource);
+    }
+
+    @Test
+    public void alreadyLocked() {
+        lack.acquire(resource);
+        thrown.expect(LackException.class);
+        lack.acquire(resource);
+    }
+
+    @Test
+    public void alreadyLockedByOther() {
+        lack.acquire(resource);
+        thrown.expect(LackException.class);
+        otherLack.acquire(resource);
+    }
+
+    @Test
+    public void alreadyReleased() {
+        lack.acquire(resource);
+        lack.release(resource);
+        thrown.expect(LackException.class);
+        otherLack.release(resource);
+    }
+
+    @Test
+    public void testReleaseAfterTtl() throws InterruptedException {
+        lack.acquire(resource);
+        Thread.sleep(1200);
+        lack.acquire(resource);
+    }
+}