Refine search
@Test public void getMultipleReadWriteSessions() { Session mockSession1 = mockSession(); Session mockSession2 = mockSession(); when(client.createSession(db)).thenReturn(mockSession1).thenReturn(mockSession2); pool = createPool(); Session session1 = pool.getReadWriteSession(); Session session2 = pool.getReadWriteSession(); verify(mockSession1).prepareReadWriteTransaction(); verify(mockSession2).prepareReadWriteTransaction(); session1.close(); session2.close(); }
Session mockSession() { Session session = mock(Session.class); when(session.getName()) .thenReturn( "projects/dummy/instances/dummy/database/dummy/sessions/session" + sessionIndex); sessionIndex++; return session; }
@Override public ReadOnlyTransaction readOnlyTransaction(TimestampBound bound) { Span span = tracer.spanBuilder(READ_ONLY_TRANSACTION).startSpan(); try (Scope s = tracer.withSpan(span)) { return pool.getReadSession().readOnlyTransaction(bound); } catch (RuntimeException e) { TraceUtil.endSpanWithFailure(span, e); throw e; } }
@Test public void getReadSessionFallsBackToWritePreparedSession() throws Exception { Session mockSession1 = mockSession(); final CountDownLatch prepareLatch = new CountDownLatch(2); doAnswer( new Answer<Void>() { @Override public Void answer(InvocationOnMock arg0) throws Throwable { prepareLatch.countDown(); return null; } }) .when(mockSession1) .prepareReadWriteTransaction(); when(client.createSession(db)).thenReturn(mockSession1); options = SessionPoolOptions.newBuilder() .setMinSessions(minSessions) .setMaxSessions(1) .setWriteSessionsFraction(1.0f) .build(); pool = createPool(); pool.getReadWriteSession().close(); prepareLatch.await(); // This session should also be write prepared. PooledSession readSession = (PooledSession) pool.getReadSession(); verify(readSession.delegate, times(2)).prepareReadWriteTransaction(); }
@Test public void keepAlive() throws Exception { options = SessionPoolOptions.newBuilder().setMinSessions(2).setMaxSessions(3).build(); Session session = mockSession(); mockKeepAlive(session); // This is cheating as we are returning the same session each but it makes the verification // easier. when(client.createSession(db)).thenReturn(session); FakeClock clock = new FakeClock(); clock.currentTimeMillis = System.currentTimeMillis(); pool = createPool(clock); Session session1 = pool.getReadSession(); Session session2 = pool.getReadSession(); session1.close(); session2.close(); runMaintainanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles); verify(session, never()).singleUse(any(TimestampBound.class)); runMaintainanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles); verify(session, times(2)).singleUse(any(TimestampBound.class)); clock.currentTimeMillis += clock.currentTimeMillis + 35 * 60 * 1000; session1 = pool.getReadSession(); session1.writeAtLeastOnce(new ArrayList<Mutation>()); session1.close(); runMaintainanceLoop(clock, pool, pool.poolMaintainer.numKeepAliveCycles); verify(session, times(3)).singleUse(any(TimestampBound.class)); pool.closeAsync().get(); }
@Test public void multiUseReadOnlyTransactionReturnsEmptyTransactionMetadata() { Transaction txnMetadata = Transaction.newBuilder().setId(ByteString.copyFromUtf8("x")).build(); PartialResultSet resultSet = PartialResultSet.newBuilder() .setMetadata(newMetadata(Type.struct(Type.StructField.of("C", Type.string())))) .build(); Mockito.when(rpc.beginTransaction(Mockito.<BeginTransactionRequest>any(), Mockito.eq(options))) .thenReturn(txnMetadata); mockRead(resultSet); ReadOnlyTransaction txn = session.readOnlyTransaction(TimestampBound.strong()); expectedException.expect(SpannerMatchers.isSpannerException(ErrorCode.INTERNAL)); txn.readRow("Dummy", Key.of(), Arrays.asList("C")); }
private void setupSession(final Session session) { ReadContext mockContext = mock(ReadContext.class); final ResultSet mockResult = mock(ResultSet.class); when(session.singleUse(any(TimestampBound.class))).thenReturn(mockContext); when(mockContext.executeQuery(any(Statement.class))) .thenAnswer( new Answer<ResultSet>() { when(mockResult.next()).thenReturn(true); doAnswer( new Answer<Void>() { .close(); doAnswer( .prepareReadWriteTransaction();
private void mockKeepAlive(Session session) { ReadContext context = mock(ReadContext.class); ResultSet resultSet = mock(ResultSet.class); when(session.singleUse(any(TimestampBound.class))).thenReturn(context); when(context.executeQuery(any(Statement.class))).thenReturn(resultSet); }
@Test public void poolClosureClosesLeakedSessions() throws Exception { Session mockSession1 = mockSession(); Session mockSession2 = mockSession(); when(client.createSession(db)).thenReturn(mockSession1).thenReturn(mockSession2); pool = createPool(); Session session1 = pool.getReadSession(); // Leaked sessions pool.getReadSession(); session1.close(); pool.closeAsync().get(); verify(mockSession1).close(); verify(mockSession2).close(); }
@Test public void createAndCloseSession() { Map<String, String> labels = new HashMap<>(); labels.put("env", "dev"); Mockito.when(spannerOptions.getSessionLabels()).thenReturn(labels); String dbName = "projects/p1/instances/i1/databases/d1"; String sessionName = dbName + "/sessions/s1"; DatabaseId db = DatabaseId.of(dbName); com.google.spanner.v1.Session sessionProto = com.google.spanner.v1.Session.newBuilder() .setName(sessionName) .putAllLabels(labels) .build(); Mockito.when(rpc.createSession(Mockito.eq(dbName), Mockito.eq(labels), options.capture())) .thenReturn(sessionProto); Session session = impl.createSession(db); assertThat(session.getName()).isEqualTo(sessionName); session.close(); // The same channelHint is passed for deleteSession (contained in "options"). Mockito.verify(rpc).deleteSession(sessionName, options.getValue()); }
@Test public void poolWorksWhenSessionNotFound() { Session mockSession1 = mockSession(); Session mockSession2 = mockSession(); doThrow(SpannerExceptionFactory.newSpannerException(ErrorCode.NOT_FOUND, "Session not found")) .when(mockSession1) .prepareReadWriteTransaction(); when(client.createSession(db)).thenReturn(mockSession1).thenReturn(mockSession2); pool = createPool(); assertThat(((PooledSession) pool.getReadWriteSession()).delegate).isEqualTo(mockSession2); }
@Test public void prepareClosesOldSingleUseContext() { ReadContext ctx = session.singleUse(TimestampBound.strong()); Mockito.when(rpc.beginTransaction(Mockito.<BeginTransactionRequest>any(), Mockito.eq(options))) .thenReturn(Transaction.newBuilder().setId(ByteString.copyFromUtf8("t1")).build()); session.prepareReadWriteTransaction(); expectedException.expect(IllegalStateException.class); expectedException.expectMessage("invalidated"); ctx.read("Dummy", KeySet.all(), Arrays.asList("C")); }
@Test public void writeClosesOldSingleUseContext() throws ParseException { ReadContext ctx = session.singleUse(TimestampBound.strong()); Mockito.when(rpc.commit(Mockito.<CommitRequest>any(), Mockito.eq(options))) .thenReturn( CommitResponse.newBuilder() .setCommitTimestamp(Timestamps.parse("2015-10-01T10:54:20.021Z")) .build()); session.writeAtLeastOnce(Arrays.<Mutation>asList()); expectedException.expect(IllegalStateException.class); expectedException.expectMessage("invalidated"); ctx.read("Dummy", KeySet.all(), Arrays.asList("C")); }
@Test public void poolClosesEvenIfPreparationFails() throws Exception { Session session = mockSession(); when(client.createSession(db)).thenReturn(session); final CountDownLatch insidePrepare = new CountDownLatch(1); final CountDownLatch releasePrepare = new CountDownLatch(1); doAnswer( new Answer<Session>() { @Override public Session answer(InvocationOnMock invocation) throws Throwable { insidePrepare.countDown(); releasePrepare.await(); throw SpannerExceptionFactory.newSpannerException(new RuntimeException()); } }) .when(session) .prepareReadWriteTransaction(); pool = createPool(); AtomicBoolean failed = new AtomicBoolean(false); CountDownLatch latch = new CountDownLatch(1); getReadWriteSessionAsync(latch, failed); insidePrepare.await(); ListenableFuture<Void> f = pool.closeAsync(); releasePrepare.countDown(); f.get(); assertThat(f.isDone()).isTrue(); }
Session session3 = mockSession(); final AtomicInteger numSessionClosed = new AtomicInteger(); when(client.createSession(db)).thenReturn(session1).thenReturn(session2).thenReturn(session3); for (Session session : new Session[] {session1, session2, session3}) { doAnswer( new Answer<Void>() { .close(); pool = createPool(clock); pool.getReadSession().close(); runMaintainanceLoop(clock, pool, pool.poolMaintainer.numClosureCycles); assertThat(numSessionClosed.get()).isEqualTo(0); Session readSession1 = pool.getReadSession(); Session readSession2 = pool.getReadSession(); Session readSession3 = pool.getReadSession(); readSession1.close(); readSession2.close(); readSession3.close(); pool.getReadSession().close(); pool.getReadSession().close(); pool.getReadSession().close(); runMaintainanceLoop(clock, pool, pool.poolMaintainer.numClosureCycles);
@Test public void writeAtLeastOnce() throws ParseException { String timestampString = "2015-10-01T10:54:20.021Z"; ArgumentCaptor<CommitRequest> commit = ArgumentCaptor.forClass(CommitRequest.class); CommitResponse response = CommitResponse.newBuilder().setCommitTimestamp(Timestamps.parse(timestampString)).build(); Mockito.when(rpc.commit(commit.capture(), Mockito.eq(options))).thenReturn(response); Timestamp timestamp = session.writeAtLeastOnce( Arrays.asList(Mutation.newInsertBuilder("T").set("C").to("x").build())); assertThat(timestamp.getSeconds()) .isEqualTo(utcTimeSeconds(2015, Calendar.OCTOBER, 1, 10, 54, 20)); assertThat(timestamp.getNanos()).isEqualTo(TimeUnit.MILLISECONDS.toNanos(21)); CommitRequest request = commit.getValue(); assertThat(request.getSingleUseTransaction()).isNotNull(); assertThat(request.getSingleUseTransaction().getReadWrite()).isNotNull(); com.google.spanner.v1.Mutation mutation = com.google.spanner.v1.Mutation.newBuilder() .setInsert( Write.newBuilder() .setTable("T") .addColumns("C") .addValues( ListValue.newBuilder() .addValues(com.google.protobuf.Value.newBuilder().setStringValue("x")))) .build(); assertThat(request.getMutationsList()).containsExactly(mutation); }
@Test public void failOnPoolExhaustion() { options = SessionPoolOptions.newBuilder() .setMinSessions(1) .setMaxSessions(1) .setFailIfPoolExhausted() .build(); Session mockSession = mockSession(); when(client.createSession(db)).thenReturn(mockSession); pool = createPool(); Session session1 = pool.getReadSession(); expectedException.expect(isSpannerException(ErrorCode.RESOURCE_EXHAUSTED)); pool.getReadSession(); session1.close(); session1 = pool.getReadSession(); assertThat(session1).isNotNull(); session1.close(); }
private void closeSession(PooledSession sess) { try { sess.delegate.close(); } catch (SpannerException e) { // Backend will delete these sessions after a while even if we fail to close them. if (logger.isLoggable(Level.FINE)) { logger.log(Level.FINE, "Failed to close session: " + sess.getName(), e); } } finally { synchronized (lock) { allSessions.remove(sess); if (isClosed()) { decrementPendingClosures(); return; } // Create a new session if needed to unblock some waiter. if (numWaiters() > numSessionsBeingCreated) { createSession(); } } } }
@Override public ReadOnlyTransaction singleUseReadOnlyTransaction() { Span span = tracer.spanBuilder(READ_ONLY_TRANSACTION).startSpan(); try (Scope s = tracer.withSpan(span)) { return pool.getReadSession().singleUseReadOnlyTransaction(); } catch (RuntimeException e) { TraceUtil.endSpanWithFailure(span, e); throw e; } }
@Override public ReadContext singleUse() { Span span = tracer.spanBuilder(READ_ONLY_TRANSACTION).startSpan(); try (Scope s = tracer.withSpan(span)) { return pool.getReadSession().singleUse(); } catch (RuntimeException e) { TraceUtil.endSpanWithFailure(span, e); throw e; } }