Skip to content

Commit

Permalink
feat: add LockHint feature (#3588)
Browse files Browse the repository at this point in the history 
* feat: add LockHint feature

* fix typo

* updated logic to make sure lockHint get propagated only for ReadWriteTransaction

* add warning message when using lock hint with unsupported transaction
  • Loading branch information
@rahul2393
rahul2393 authored Jan 9, 2025
1 parent 7af512b commit 326442b
Show file tree
Hide file tree
Showing 5 changed files with 163 additions and 3 deletions.
11 changes: 11 additions & 0 deletions google-cloud-spanner/src/main/java/com/google/cloud/spanner/AbstractReadContext.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import java.util.Map;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicLong;
import java.util.logging.Logger;
import javax.annotation.Nullable;
import javax.annotation.concurrent.GuardedBy;

Expand All @@ -67,6 +68,7 @@
*/
abstract class AbstractReadContext
implements ReadContext, AbstractResultSet.Listener, SessionTransaction {
private static final Logger logger = Logger.getLogger(AbstractReadContext.class.getName());

abstract static class Builder<B extends Builder<?, T>, T extends AbstractReadContext> {
private SessionImpl session;
Expand Down Expand Up @@ -951,6 +953,15 @@ ResultSet readInternalWithOptions(
} else if (defaultDirectedReadOptions != null) {
builder.setDirectedReadOptions(defaultDirectedReadOptions);
}
if (readOptions.hasLockHint()) {
if (isReadOnly()) {
logger.warning(
"Lock hint is only supported for ReadWrite transactions. "
+ "Overriding lock hint to default unspecified.");
} else {
builder.setLockHint(readOptions.lockHint());
}
}
final int prefetchChunks =
readOptions.hasPrefetchChunks() ? readOptions.prefetchChunks() : defaultPrefetchChunks;
ResumableStreamIterator stream =
Expand Down
55 changes: 54 additions & 1 deletion google-cloud-spanner/src/main/java/com/google/cloud/spanner/Options.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@

import com.google.common.base.Preconditions;
import com.google.spanner.v1.DirectedReadOptions;
import com.google.spanner.v1.ReadRequest.LockHint;
import com.google.spanner.v1.ReadRequest.OrderBy;
import com.google.spanner.v1.RequestOptions.Priority;
import java.io.Serializable;
Expand Down Expand Up @@ -75,6 +76,25 @@ public static RpcOrderBy fromProto(OrderBy proto) {
}
}

public enum RpcLockHint {
UNSPECIFIED(LockHint.LOCK_HINT_UNSPECIFIED),
SHARED(LockHint.LOCK_HINT_SHARED),
EXCLUSIVE(LockHint.LOCK_HINT_EXCLUSIVE);

private final LockHint proto;

RpcLockHint(LockHint proto) {
this.proto = Preconditions.checkNotNull(proto);
}

public static RpcLockHint fromProto(LockHint proto) {
for (RpcLockHint e : RpcLockHint.values()) {
if (e.proto.equals(proto)) return e;
}
return RpcLockHint.UNSPECIFIED;
}
}

/** Marker interface to mark options applicable to both Read and Query operations */
public interface ReadAndQueryOption extends ReadOption, QueryOption {}

Expand Down Expand Up @@ -160,6 +180,10 @@ public static ReadOption orderBy(RpcOrderBy orderBy) {
return new OrderByOption(orderBy);
}

public static ReadOption lockHint(RpcLockHint orderBy) {
return new LockHintOption(orderBy);
}

/**
* Specifying this will allow the client to prefetch up to {@code prefetchChunks} {@code
* PartialResultSet} chunks for read and query. The data size of each chunk depends on the server
Expand Down Expand Up @@ -469,6 +493,7 @@ void appendToOptions(Options options) {
private DirectedReadOptions directedReadOptions;
private DecodeMode decodeMode;
private RpcOrderBy orderBy;
private RpcLockHint lockHint;

// Construction is via factory methods below.
private Options() {}
Expand Down Expand Up @@ -605,6 +630,14 @@ OrderBy orderBy() {
return orderBy == null ? null : orderBy.proto;
}

boolean hasLockHint() {
return lockHint != null;
}

LockHint lockHint() {
return lockHint == null ? null : lockHint.proto;
}

@Override
public String toString() {
StringBuilder b = new StringBuilder();
Expand Down Expand Up @@ -661,6 +694,9 @@ public String toString() {
if (orderBy != null) {
b.append("orderBy: ").append(orderBy).append(' ');
}
if (lockHint != null) {
b.append("lockHint: ").append(lockHint).append(' ');
}
return b.toString();
}

Expand Down Expand Up @@ -700,7 +736,8 @@ public boolean equals(Object o) {
&& Objects.equals(withExcludeTxnFromChangeStreams(), that.withExcludeTxnFromChangeStreams())
&& Objects.equals(dataBoostEnabled(), that.dataBoostEnabled())
&& Objects.equals(directedReadOptions(), that.directedReadOptions())
&& Objects.equals(orderBy(), that.orderBy());
&& Objects.equals(orderBy(), that.orderBy())
&& Objects.equals(lockHint(), that.lockHint());
}

@Override
Expand Down Expand Up @@ -760,6 +797,9 @@ public int hashCode() {
if (orderBy != null) {
result = 31 * result + orderBy.hashCode();
}
if (lockHint != null) {
result = 31 * result + lockHint.hashCode();
}
return result;
}

Expand Down Expand Up @@ -853,6 +893,19 @@ void appendToOptions(Options options) {
}
}

static class LockHintOption extends InternalOption implements ReadOption {
private final RpcLockHint lockHint;

LockHintOption(RpcLockHint lockHint) {
this.lockHint = lockHint;
}

@Override
void appendToOptions(Options options) {
options.lockHint = lockHint;
}
}

static final class DataBoostQueryOption extends InternalOption implements ReadAndQueryOption {

private final Boolean dataBoostEnabled;
Expand Down
16 changes: 16 additions & 0 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/AbstractReadContextTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,7 @@
import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
import com.google.spanner.v1.ReadRequest;
import com.google.spanner.v1.ReadRequest.LockHint;
import com.google.spanner.v1.ReadRequest.OrderBy;
import com.google.spanner.v1.RequestOptions;
import com.google.spanner.v1.RequestOptions.Priority;
Expand Down Expand Up @@ -241,6 +242,21 @@ public void testGetReadRequestBuilderWithOrderBy() {
assertEquals(OrderBy.ORDER_BY_NO_ORDER, request.getOrderBy());
}

@Test
public void testGetReadRequestBuilderWithLockHint() {
ReadRequest request =
ReadRequest.newBuilder()
.setSession(
SessionName.of("[PROJECT]", "[INSTANCE]", "[DATABASE]", "[SESSION]").toString())
.setTransaction(TransactionSelector.newBuilder().build())
.setTable("table110115790")
.setIndex("index100346066")
.addAllColumns(new ArrayList<String>())
.setLockHintValue(2)
.build();
assertEquals(LockHint.LOCK_HINT_EXCLUSIVE, request.getLockHint());
}

@Test
public void testGetExecuteBatchDmlRequestBuilderWithPriority() {
ExecuteBatchDmlRequest.Builder request =
Expand Down
49 changes: 49 additions & 0 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/DatabaseClientImplTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,7 @@
import com.google.cloud.spanner.AsyncTransactionManager.TransactionContextFuture;
import com.google.cloud.spanner.MockSpannerServiceImpl.SimulatedExecutionTime;
import com.google.cloud.spanner.MockSpannerServiceImpl.StatementResult;
import com.google.cloud.spanner.Options.RpcLockHint;
import com.google.cloud.spanner.Options.RpcOrderBy;
import com.google.cloud.spanner.Options.RpcPriority;
import com.google.cloud.spanner.Options.TransactionOption;
Expand Down Expand Up @@ -90,6 +91,7 @@
import com.google.spanner.v1.ExecuteSqlRequest.QueryMode;
import com.google.spanner.v1.ExecuteSqlRequest.QueryOptions;
import com.google.spanner.v1.ReadRequest;
import com.google.spanner.v1.ReadRequest.LockHint;
import com.google.spanner.v1.ReadRequest.OrderBy;
import com.google.spanner.v1.RequestOptions.Priority;
import com.google.spanner.v1.ResultSetMetadata;
Expand Down Expand Up @@ -1754,6 +1756,53 @@ public void testExecuteReadWithOrderByOption() {
assertEquals(OrderBy.ORDER_BY_NO_ORDER, request.getOrderBy());
}

@Test
public void testUnsupportedTransactionWithLockHintOption() {
DatabaseClient client =
spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));
try (ResultSet resultSet =
client
.singleUse()
.read(
READ_TABLE_NAME,
KeySet.singleKey(Key.of(1L)),
READ_COLUMN_NAMES,
Options.lockHint(RpcLockHint.EXCLUSIVE))) {
consumeResults(resultSet);
}

List<ReadRequest> requests = mockSpanner.getRequestsOfType(ReadRequest.class);
assertThat(requests).hasSize(1);
ReadRequest request = requests.get(0);
// lock hint is only supported in ReadWriteTransaction
assertEquals(LockHint.LOCK_HINT_UNSPECIFIED, request.getLockHint());
}

@Test
public void testReadWriteTransactionWithLockHint() {
DatabaseClient client =
spanner.getDatabaseClient(DatabaseId.of(TEST_PROJECT, TEST_INSTANCE, TEST_DATABASE));

TransactionRunner runner = client.readWriteTransaction();
runner.run(
transaction -> {
try (ResultSet resultSet =
transaction.read(
READ_TABLE_NAME,
KeySet.singleKey(Key.of(1L)),
READ_COLUMN_NAMES,
Options.lockHint(RpcLockHint.EXCLUSIVE))) {
consumeResults(resultSet);
}
return null;
});

List<ReadRequest> requests = mockSpanner.getRequestsOfType(ReadRequest.class);
assertThat(requests).hasSize(1);
ReadRequest request = requests.get(0);
assertEquals(LockHint.LOCK_HINT_EXCLUSIVE, request.getLockHint());
}

@Test
public void testExecuteReadWithDirectedReadOptions() {
DatabaseClient client =
Expand Down
35 changes: 33 additions & 2 deletions google-cloud-spanner/src/test/java/com/google/cloud/spanner/OptionsTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,13 @@
import static org.junit.Assert.assertThrows;
import static org.junit.Assert.assertTrue;

import com.google.cloud.spanner.Options.RpcLockHint;
import com.google.cloud.spanner.Options.RpcOrderBy;
import com.google.cloud.spanner.Options.RpcPriority;
import com.google.spanner.v1.DirectedReadOptions;
import com.google.spanner.v1.DirectedReadOptions.IncludeReplicas;
import com.google.spanner.v1.DirectedReadOptions.ReplicaSelection;
import com.google.spanner.v1.ReadRequest.LockHint;
import com.google.spanner.v1.ReadRequest.OrderBy;
import com.google.spanner.v1.RequestOptions.Priority;
import org.junit.Test;
Expand Down Expand Up @@ -83,7 +85,8 @@ public void allOptionsPresent() {
Options.prefetchChunks(1),
Options.dataBoostEnabled(true),
Options.directedRead(DIRECTED_READ_OPTIONS),
Options.orderBy(RpcOrderBy.NO_ORDER));
Options.orderBy(RpcOrderBy.NO_ORDER),
Options.lockHint(Options.RpcLockHint.SHARED));
assertThat(options.hasLimit()).isTrue();
assertThat(options.limit()).isEqualTo(10);
assertThat(options.hasPrefetchChunks()).isTrue();
Expand All @@ -92,6 +95,7 @@ public void allOptionsPresent() {
assertTrue(options.dataBoostEnabled());
assertTrue(options.hasDirectedReadOptions());
assertTrue(options.hasOrderBy());
assertTrue(options.hasLockHint());
assertEquals(DIRECTED_READ_OPTIONS, options.directedReadOptions());
}

Expand All @@ -107,6 +111,7 @@ public void allOptionsAbsent() {
assertThat(options.hasDataBoostEnabled()).isFalse();
assertThat(options.hasDirectedReadOptions()).isFalse();
assertThat(options.hasOrderBy()).isFalse();
assertThat(options.hasLockHint()).isFalse();
assertNull(options.withExcludeTxnFromChangeStreams());
assertThat(options.toString()).isEqualTo("");
assertThat(options.equals(options)).isTrue();
Expand Down Expand Up @@ -189,7 +194,8 @@ public void readOptionsTest() {
Options.tag(tag),
Options.dataBoostEnabled(true),
Options.directedRead(DIRECTED_READ_OPTIONS),
Options.orderBy(RpcOrderBy.NO_ORDER));
Options.orderBy(RpcOrderBy.NO_ORDER),
Options.lockHint(RpcLockHint.SHARED));

assertThat(options.toString())
.isEqualTo(
Expand All @@ -207,11 +213,15 @@ public void readOptionsTest() {
+ " "
+ "orderBy: "
+ RpcOrderBy.NO_ORDER
+ " "
+ "lockHint: "
+ RpcLockHint.SHARED
+ " ");
assertThat(options.tag()).isEqualTo(tag);
assertEquals(dataBoost, options.dataBoostEnabled());
assertEquals(DIRECTED_READ_OPTIONS, options.directedReadOptions());
assertEquals(OrderBy.ORDER_BY_NO_ORDER, options.orderBy());
assertEquals(LockHint.LOCK_HINT_SHARED, options.lockHint());
}

@Test
Expand Down Expand Up @@ -373,6 +383,14 @@ public void testReadOptionsOrderBy() {
assertEquals("orderBy: " + orderBy + " ", options.toString());
}

@Test
public void testReadOptionsLockHint() {
RpcLockHint lockHint = RpcLockHint.SHARED;
Options options = Options.fromReadOptions(Options.lockHint(lockHint));
assertTrue(options.hasLockHint());
assertEquals("lockHint: " + lockHint + " ", options.toString());
}

@Test
public void testReadOptionsWithOrderByEquality() {
Options optionsWithNoOrderBy1 = Options.fromReadOptions(Options.orderBy(RpcOrderBy.NO_ORDER));
Expand All @@ -383,6 +401,19 @@ public void testReadOptionsWithOrderByEquality() {
assertFalse(optionsWithNoOrderBy1.equals(optionsWithPkOrder));
}

@Test
public void testReadOptionsWithLockHintEquality() {
Options optionsWithSharedLockHint1 =
Options.fromReadOptions(Options.lockHint(RpcLockHint.SHARED));
Options optionsWithSharedLockHint2 =
Options.fromReadOptions(Options.lockHint(RpcLockHint.SHARED));
assertEquals(optionsWithSharedLockHint1, optionsWithSharedLockHint2);

Options optionsWithExclusiveLock =
Options.fromReadOptions(Options.lockHint(RpcLockHint.EXCLUSIVE));
assertNotEquals(optionsWithSharedLockHint1, optionsWithExclusiveLock);
}

@Test
public void testQueryOptionsPriority() {
RpcPriority priority = RpcPriority.MEDIUM;
Expand Down

0 comments on commit 326442b

Please sign in to comment.