@Test public void unlocksWhenFailedToRegisterLockAfterStartingTransaction() { when(response.immutableTimestamp()).thenReturn(immutableTimestampResponse); when(delegate.startIdentifiedAtlasDbTransactionBatch(1)).thenReturn(ImmutableList.of(response)); ImmutableSet<LockToken> locks = ImmutableSet.of(immutableTsLock); doThrow(new RuntimeException()).when(refresher).registerLocks(locks); assertThatThrownBy(() -> timelock.startIdentifiedAtlasDbTransactionBatch(1)) .isInstanceOf(RuntimeException.class); verify(refresher).unregisterLocks(locks); }
@Override public List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTransactionBatch(int count) { List<StartIdentifiedAtlasDbTransactionResponse> responses = executeOnTimeLock( () -> delegate.startIdentifiedAtlasDbTransactionBatch(count)); Set<LockToken> immutableTsLocks = responses .stream() .map(response -> response.immutableTimestamp().getLock()) .collect(Collectors.toSet()); try { lockRefresher.registerLocks(immutableTsLocks); } catch (Throwable t) { unlock(immutableTsLocks); throw Throwables.throwUncheckedException(t); } return responses; }
TimeLockClient implements AutoCloseable, TimelockService { @Override public List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTransactionBatch(int count) { List<StartIdentifiedAtlasDbTransactionResponse> responses = executeOnTimeLock( () -> delegate.startIdentifiedAtlasDbTransactionBatch(count)); Set<LockToken> immutableTsLocks = responses .stream() .map(response -> response.immutableTimestamp().getLock()) .collect(Collectors.toSet()); try { lockRefresher.registerLocks(immutableTsLocks); } catch (Throwable t) { unlock(immutableTsLocks); throw Throwables.throwUncheckedException(t); } return responses; } }
TimeLockClient implements AutoCloseable, TimelockService { @Override public List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTransactionBatch(int count) { List<StartIdentifiedAtlasDbTransactionResponse> responses = executeOnTimeLock( () -> delegate.startIdentifiedAtlasDbTransactionBatch(count)); Set<LockToken> immutableTsLocks = responses .stream() .map(response -> response.immutableTimestamp().getLock()) .collect(Collectors.toSet()); try { lockRefresher.registerLocks(immutableTsLocks); } catch (Throwable t) { unlock(immutableTsLocks); throw Throwables.throwUncheckedException(t); } return responses; } @VisibleForTesting  TimeLockClient(TimelockService delegate, CloseableTimestampService timestampService,
            LockRefresher lockRefresher, TimeLockUnlocker unlocker); }
TimeLockClient implements AutoCloseable, TimelockService { @Override public List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTransactionBatch(int count) { List<StartIdentifiedAtlasDbTransactionResponse> responses = executeOnTimeLock( () -> delegate.startIdentifiedAtlasDbTransactionBatch(count)); Set<LockToken> immutableTsLocks = responses .stream() .map(response -> response.immutableTimestamp().getLock()) .collect(Collectors.toSet()); try { lockRefresher.registerLocks(immutableTsLocks); } catch (Throwable t) { unlock(immutableTsLocks); throw Throwables.throwUncheckedException(t); } return responses; } @VisibleForTesting  TimeLockClient(TimelockService delegate, CloseableTimestampService timestampService,
            LockRefresher lockRefresher, TimeLockUnlocker unlocker); static TimeLockClient createDefault(TimelockService timelockService); static TimeLockClient withSynchronousUnlocker(TimelockService timelockService); @Override boolean isInitialized(); @Override long getFreshTimestamp(); @Override long getCommitTimestamp(long startTs, LockToken commitLocksToken); @Override TimestampRange getFreshTimestamps(int numTimestampsRequested); @Override LockImmutableTimestampResponse lockImmutableTimestamp(); @Override List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTransactionBatch(int count); @Override long getImmutableTimestamp(); @Override LockResponse lock(LockRequest request); @Override WaitForLocksResponse waitForLocks(WaitForLocksRequest request); @Override Set<LockToken> refreshLockLeases(Set<LockToken> tokens); @Override Set<LockToken> unlock(Set<LockToken> tokens); @Override void tryUnlock(Set<LockToken> tokens); @Override long currentTimeMillis(); @Override void close(); }
TimeLockClient implements AutoCloseable, TimelockService { @Override public List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTransactionBatch(int count) { List<StartIdentifiedAtlasDbTransactionResponse> responses = executeOnTimeLock( () -> delegate.startIdentifiedAtlasDbTransactionBatch(count)); Set<LockToken> immutableTsLocks = responses .stream() .map(response -> response.immutableTimestamp().getLock()) .collect(Collectors.toSet()); try { lockRefresher.registerLocks(immutableTsLocks); } catch (Throwable t) { unlock(immutableTsLocks); throw Throwables.throwUncheckedException(t); } return responses; } @VisibleForTesting  TimeLockClient(TimelockService delegate, CloseableTimestampService timestampService,
            LockRefresher lockRefresher, TimeLockUnlocker unlocker); static TimeLockClient createDefault(TimelockService timelockService); static TimeLockClient withSynchronousUnlocker(TimelockService timelockService); @Override boolean isInitialized(); @Override long getFreshTimestamp(); @Override long getCommitTimestamp(long startTs, LockToken commitLocksToken); @Override TimestampRange getFreshTimestamps(int numTimestampsRequested); @Override LockImmutableTimestampResponse lockImmutableTimestamp(); @Override List<StartIdentifiedAtlasDbTransactionResponse> startIdentifiedAtlasDbTransactionBatch(int count); @Override long getImmutableTimestamp(); @Override LockResponse lock(LockRequest request); @Override WaitForLocksResponse waitForLocks(WaitForLocksRequest request); @Override Set<LockToken> refreshLockLeases(Set<LockToken> tokens); @Override Set<LockToken> unlock(Set<LockToken> tokens); @Override void tryUnlock(Set<LockToken> tokens); @Override long currentTimeMillis(); @Override void close();  }