diff --git a/java/google/registry/model/ofy/ObjectifyService.java b/java/google/registry/model/ofy/ObjectifyService.java index 8b44e3309..f88c738e0 100644 --- a/java/google/registry/model/ofy/ObjectifyService.java +++ b/java/google/registry/model/ofy/ObjectifyService.java @@ -109,7 +109,7 @@ public class ObjectifyService { // examine the number of requests sent to datastore. AsyncDatastoreService service = super.createRawAsyncDatastoreService(cfg); return RegistryEnvironment.get().equals(RegistryEnvironment.UNITTEST) - ? new RequestCountingAsyncDatastoreService(service) + ? new RequestCapturingAsyncDatastoreService(service) : service; }}); diff --git a/java/google/registry/model/ofy/RequestCountingAsyncDatastoreService.java b/java/google/registry/model/ofy/RequestCapturingAsyncDatastoreService.java similarity index 73% rename from java/google/registry/model/ofy/RequestCountingAsyncDatastoreService.java rename to java/google/registry/model/ofy/RequestCapturingAsyncDatastoreService.java index 03bdf285e..6ebcf2081 100644 --- a/java/google/registry/model/ofy/RequestCountingAsyncDatastoreService.java +++ b/java/google/registry/model/ofy/RequestCapturingAsyncDatastoreService.java @@ -14,6 +14,8 @@ package google.registry.model.ofy; +import static java.util.Collections.synchronizedList; + import com.google.appengine.api.datastore.AsyncDatastoreService; import com.google.appengine.api.datastore.DatastoreAttributes; import com.google.appengine.api.datastore.Entity; @@ -25,39 +27,41 @@ import com.google.appengine.api.datastore.PreparedQuery; import com.google.appengine.api.datastore.Query; import com.google.appengine.api.datastore.Transaction; import com.google.appengine.api.datastore.TransactionOptions; +import com.google.common.collect.ImmutableList; +import java.util.ArrayList; import java.util.Collection; import java.util.List; import java.util.Map; import java.util.concurrent.Future; -import java.util.concurrent.atomic.AtomicInteger; /** A proxy for {@link AsyncDatastoreService} that exposes call counts. */ -public class RequestCountingAsyncDatastoreService implements AsyncDatastoreService { +public class RequestCapturingAsyncDatastoreService implements AsyncDatastoreService { private final AsyncDatastoreService delegate; - // We use static counters because we care about overall calls to datastore, not calls via a - // specific instance of the service. + // Each outer lists represents datastore operations, with inner lists representing the keys or + // entities involved in that operation. We use static lists because we care about overall calls to + // datastore, not calls via a specific instance of the service. - private static AtomicInteger reads = new AtomicInteger(); - private static AtomicInteger puts = new AtomicInteger(); - private static AtomicInteger deletes = new AtomicInteger(); + private static List> reads = synchronizedList(new ArrayList>()); + private static List> deletes = synchronizedList(new ArrayList>()); + private static List> puts = synchronizedList(new ArrayList>()); - RequestCountingAsyncDatastoreService(AsyncDatastoreService delegate) { + RequestCapturingAsyncDatastoreService(AsyncDatastoreService delegate) { this.delegate = delegate; } - public static int getReadsCount() { - return reads.get(); + public static List> getReads() { + return reads; } - public static int getPutsCount() { - return puts.get(); + public static List> getDeletes() { + return deletes; } - public static int getDeletesCount() { - return deletes.get(); + public static List> getPuts() { + return puts; } @Override @@ -107,49 +111,49 @@ public class RequestCountingAsyncDatastoreService implements AsyncDatastoreServi @Override public Future delete(Key... keys) { - deletes.incrementAndGet(); + deletes.add(ImmutableList.copyOf(keys)); return delegate.delete(keys); } @Override public Future delete(Iterable keys) { - deletes.incrementAndGet(); + deletes.add(ImmutableList.copyOf(keys)); return delegate.delete(keys); } @Override public Future delete(Transaction transaction, Key... keys) { - deletes.incrementAndGet(); + deletes.add(ImmutableList.copyOf(keys)); return delegate.delete(transaction, keys); } @Override public Future delete(Transaction transaction, Iterable keys) { - deletes.incrementAndGet(); + deletes.add(ImmutableList.copyOf(keys)); return delegate.delete(transaction, keys); } @Override public Future get(Key key) { - reads.incrementAndGet(); + reads.add(ImmutableList.of(key)); return delegate.get(key); } @Override public Future> get(Iterable keys) { - reads.incrementAndGet(); + reads.add(ImmutableList.copyOf(keys)); return delegate.get(keys); } @Override public Future get(Transaction transaction, Key key) { - reads.incrementAndGet(); + reads.add(ImmutableList.of(key)); return delegate.get(transaction, key); } @Override public Future> get(Transaction transaction, Iterable keys) { - reads.incrementAndGet(); + reads.add(ImmutableList.copyOf(keys)); return delegate.get(transaction, keys); } @@ -165,25 +169,25 @@ public class RequestCountingAsyncDatastoreService implements AsyncDatastoreServi @Override public Future put(Entity entity) { - puts.incrementAndGet(); + puts.add(ImmutableList.of(entity)); return delegate.put(entity); } @Override public Future> put(Iterable entities) { - puts.incrementAndGet(); + puts.add(ImmutableList.copyOf(entities)); return delegate.put(entities); } @Override public Future put(Transaction transaction, Entity entity) { - puts.incrementAndGet(); + puts.add(ImmutableList.of(entity)); return delegate.put(transaction, entity); } @Override public Future> put(Transaction transaction, Iterable entities) { - puts.incrementAndGet(); + puts.add(ImmutableList.copyOf(entities)); return delegate.put(transaction, entities); } } diff --git a/javatests/google/registry/model/domain/DomainResourceTest.java b/javatests/google/registry/model/domain/DomainResourceTest.java index 0d3190da7..1ef72f2d3 100644 --- a/javatests/google/registry/model/domain/DomainResourceTest.java +++ b/javatests/google/registry/model/domain/DomainResourceTest.java @@ -16,6 +16,7 @@ package google.registry.model.domain; import static com.google.appengine.tools.development.testing.LocalMemcacheServiceTestConfig.getLocalMemcacheService; import static com.google.common.collect.Iterables.getOnlyElement; +import static com.google.common.collect.Iterables.skip; import static com.google.common.truth.Truth.assertThat; import static google.registry.model.EppResourceUtils.loadByUniqueId; import static google.registry.testing.DatastoreHelper.cloneAndSetAutoTimestamps; @@ -54,7 +55,7 @@ import google.registry.model.eppoutput.Response; import google.registry.model.eppoutput.Result; import google.registry.model.eppoutput.Result.Code; import google.registry.model.host.HostResource; -import google.registry.model.ofy.RequestCountingAsyncDatastoreService; +import google.registry.model.ofy.RequestCapturingAsyncDatastoreService; import google.registry.model.poll.PollMessage; import google.registry.model.registry.Registry; import google.registry.model.reporting.HistoryEntry; @@ -432,7 +433,7 @@ public class DomainResourceTest extends EntityTestCase { // Clear out memcache so that we count actual datastore calls. getLocalMemcacheService().flushAll( new LocalRpcService.Status(), MemcacheFlushRequest.newBuilder().build()); - int previousReads = RequestCountingAsyncDatastoreService.getReadsCount(); + int numPreviousReads = RequestCapturingAsyncDatastoreService.getReads().size(); EppXmlTransformer.marshal( EppOutput.create(new Response.Builder() .setResult(Result.create(Code.Success)) @@ -440,6 +441,7 @@ public class DomainResourceTest extends EntityTestCase { .setTrid(Trid.create(null, "abc")) .build()), ValidationMode.STRICT); - assertThat(RequestCountingAsyncDatastoreService.getReadsCount() - previousReads).isEqualTo(1); + // Assert that there was only one call to datastore (that may have loaded many keys). + assertThat(skip(RequestCapturingAsyncDatastoreService.getReads(), numPreviousReads)).hasSize(1); } }