Load foreign keys more efficiently for xml marshalling.

Before this CL, each contact and host was independently
loaded via the ReferenceUnion adapter. Since fields are
processed serially by JAXB, this means worst-case there
were 17 loads, best case 3 (the 3 required contacts) and
usual case 5-6 (some hosts). This CL reduces that to 1
datastore roundtrip in all cases.

A side effect of this CL is the further hollowing-out of
ReferenceUnion, since it no longer is involved in
marshalling at all.
-------------
Created by MOE: https://github.com/google/moe
MOE_MIGRATED_REVID=123712842

java/google/registry/model/domain/DomainApplication.java

@@ -43,9 +43,9 @@ import javax.xml.bind.annotation.XmlType;
     "fullyQualifiedDomainName",
     "repoId",
     "status",
-    "registrant",
+    "marshalledRegistrant",
-    "contacts",
+    "marshalledContacts",
-    "nameservers",
+    "marshalledNameservers",
     "currentSponsorClientId",
     "creationClientId",
     "creationTime",

java/google/registry/model/domain/DomainBase.java

@@ -14,11 +14,15 @@
 
 package google.registry.model.domain;
 
+import static com.google.common.base.Preconditions.checkArgument;
+import static com.google.common.base.Predicates.not;
 import static com.google.common.base.Strings.emptyToNull;
+import static com.google.common.collect.Iterables.all;
+import static com.google.common.collect.Iterables.any;
 import static com.google.common.collect.Sets.difference;
 import static com.google.common.collect.Sets.union;
-import static google.registry.model.domain.DesignatedContact.Type.REGISTRANT;
 import static google.registry.model.ofy.ObjectifyService.ofy;
+import static google.registry.util.CollectionUtils.forceEmptyToNull;
 import static google.registry.util.CollectionUtils.nullToEmpty;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableCopy;
 import static google.registry.util.CollectionUtils.nullToEmptyImmutableSortedCopy;
@@ -27,20 +31,20 @@ import static google.registry.util.DomainNameUtils.getTldFromDomainName;
 import static google.registry.util.PreconditionsUtils.checkArgumentNotNull;
 
 import com.google.common.base.Function;
+import com.google.common.base.Predicate;
 import com.google.common.collect.FluentIterable;
 import com.google.common.collect.ImmutableSet;
 import com.google.common.collect.ImmutableSortedSet;
 
 import com.googlecode.objectify.Ref;
 import com.googlecode.objectify.annotation.Entity;
-import com.googlecode.objectify.annotation.Ignore;
 import com.googlecode.objectify.annotation.IgnoreSave;
 import com.googlecode.objectify.annotation.Index;
-import com.googlecode.objectify.annotation.OnLoad;
 import com.googlecode.objectify.condition.IfNull;
 
 import google.registry.model.EppResource;
 import google.registry.model.contact.ContactResource;
+import google.registry.model.domain.DesignatedContact.Type;
 import google.registry.model.domain.launch.LaunchNotice;
 import google.registry.model.domain.secdns.DelegationSignerData;
 import google.registry.model.eppcommon.AuthInfo;
@@ -76,40 +80,18 @@ public abstract class DomainBase extends EppResource {
   String tld;
 
   /** References to hosts that are the nameservers for the domain. */
-  @XmlElementWrapper(name = "ns")
+  @XmlTransient
-  @XmlElement(name = "hostObj")
   //TODO(b/28713909): Make this a Set<Ref<HostResource>>.
   Set<ReferenceUnion<HostResource>> nameservers;
 
   /**
-   * Associated contacts for the domain (other than registrant).
+   * The union of the contacts visible via {@link #getContacts} and {@link #getRegistrant}.
    *
-   * <p>This field is marked with {@literal @}Ignore so that {@link DomainBase} subclasses won't
+   * <p>These are stored in one field so that we can query across all contacts at once.
-   * persist it. Instead, the data in this field and in the {@link #registrant} are both stored in
-   * {@link DomainBase#allContacts} to allow for more efficient queries.
-   */
-  @Ignore
-  @XmlElement(name = "contact")
-  Set<DesignatedContact> contacts;
-
-  /**
-   * The union of the contacts in {@link #contacts} and {@link #registrant}. This is so we can query
-   * across all contacts at once. It is maintained by the builder and by {@link #unpackageContacts}.
    */
   @XmlTransient
   Set<DesignatedContact> allContacts;
 
-  /**
-   * A reference to the registrant who registered this domain.
-   *
-   * <p>This field is marked with {@literal @}Ignore so that {@link DomainBase} subclasses won't
-   * persist it. Instead, the data in this field and in the {@link DomainBase#contacts} are both
-   * stored in {@link DomainBase#allContacts} to allow for more efficient queries.
-   */
-  @Ignore
-  //TODO(b/28713909): Make this a Ref<ContactResource>.
-  ReferenceUnion<ContactResource> registrant;
-
   /** Authorization info (aka transfer secret) of the domain. */
   DomainAuthInfo authInfo;
 
@@ -139,6 +121,49 @@ public abstract class DomainBase extends EppResource {
   @XmlTransient
   String idnTableName;
 
+  /**
+   * Synchronously load all referenced contacts and hosts into the Objectify session cache.
+   *
+   * <p>This saves an extra datastore roundtrip on marshalling, since contacts, hosts, and the
+   * registrant will all be in the session cache when their respective methods are called.
+   */
+  private void preMarshal() {
+    // Calling values() blocks until loading is done.
+    ofy().load().values(union(getNameservers(), getReferencedContacts())).values();
+  }
+
+  /** JAXB java beans property to marshal nameserver hostnames. */
+  @XmlElementWrapper(name = "ns")
+  @XmlElement(name = "hostObj")
+  private ImmutableSet<String> getMarshalledNameservers() {
+    preMarshal();
+    // If there are no nameservers we must return null, or an empty "ns" element will be marshalled.
+    return forceEmptyToNull(loadNameserverFullyQualifiedHostNames());
+  }
+
+  /** JAXB java beans property to marshal non-registrant contact ids. */
+  @XmlElement(name = "contact")
+  private ImmutableSet<ForeignKeyedDesignatedContact> getMarshalledContacts() {
+    preMarshal();
+    return FluentIterable.from(getContacts())
+        .transform(
+            new Function<DesignatedContact, ForeignKeyedDesignatedContact>() {
+              @Override
+              public ForeignKeyedDesignatedContact apply(DesignatedContact designated) {
+                return ForeignKeyedDesignatedContact.create(
+                    designated.getType(),
+                    designated.getContactRef().get().getContactId());
+              }})
+        .toSet();
+  }
+
+  /** JAXB java beans property to marshal nameserver hostnames. */
+  @XmlElement(name = "registrant")
+  private String getMarshalledRegistrant() {
+    preMarshal();
+    return getRegistrant().get().getContactId();
+  }
+
   public String getFullyQualifiedDomainName() {
     return fullyQualifiedDomainName;
   }
@@ -175,12 +200,21 @@ public abstract class DomainBase extends EppResource {
         .toSet();
   }
 
+  /** A reference to the registrant who registered this domain. */
   public Ref<ContactResource> getRegistrant() {
-    return registrant.getLinked();
+    return FluentIterable
+        .from(nullToEmpty(allContacts))
+        .filter(IS_REGISTRANT)
+        .getOnlyElement()
+        .getContactRef();
   }
 
+  /** Associated contacts for the domain (other than registrant). */
   public ImmutableSet<DesignatedContact> getContacts() {
-    return nullToEmptyImmutableCopy(contacts);
+    return FluentIterable
+        .from(nullToEmpty(allContacts))
+        .filter(not(IS_REGISTRANT))
+        .toSet();
   }
 
   public AuthInfo getAuthInfo() {
@@ -201,22 +235,13 @@ public abstract class DomainBase extends EppResource {
     return tld;
   }
 
-  /**
+  /** Predicate to determine if a given {@link DesignatedContact} is the registrant. */
-   * On load, unpackage the {@link #allContacts} field, placing the registrant into
+  private static final Predicate<DesignatedContact> IS_REGISTRANT =
-   * {@link #registrant} field and all other contacts into {@link #contacts}.
+      new Predicate<DesignatedContact>() {
-   */
+        @Override
-  @OnLoad
+        public boolean apply(DesignatedContact contact) {
-  void unpackageContacts() {
+          return DesignatedContact.Type.REGISTRANT.equals(contact.type);
-    ImmutableSet.Builder<DesignatedContact> contactsBuilder = new ImmutableSet.Builder<>();
+        }};
-    for (DesignatedContact contact : nullToEmptyImmutableCopy(allContacts)) {
-      if (REGISTRANT.equals(contact.getType())){
-        registrant = ReferenceUnion.create(contact.getContactRef());
-      } else {
-        contactsBuilder.add(contact);
-      }
-    }
-    contacts = contactsBuilder.build();
-  }
 
   /** An override of {@link EppResource#asBuilder} with tighter typing. */
   @Override
@@ -225,6 +250,7 @@ public abstract class DomainBase extends EppResource {
   /** A builder for constructing {@link DomainBase}, since it is immutable. */
   public abstract static class Builder<T extends DomainBase, B extends Builder<?, ?>>
       extends EppResource.Builder<T, B> {
+
     protected Builder() {}
 
     protected Builder(T instance) {
@@ -236,11 +262,8 @@ public abstract class DomainBase extends EppResource {
       T instance = getInstance();
       checkArgumentNotNull(
           emptyToNull(instance.fullyQualifiedDomainName), "Missing fullyQualifiedDomainName");
-      checkArgumentNotNull(instance.registrant, "Missing registrant");
+      checkArgument(any(instance.allContacts, IS_REGISTRANT), "Missing registrant");
       instance.tld = getTldFromDomainName(instance.fullyQualifiedDomainName);
-      instance.allContacts = union(
-              instance.getContacts(),
-              DesignatedContact.create(REGISTRANT, instance.registrant.getLinked()));
       return super.build();
     }
 
@@ -255,7 +278,10 @@ public abstract class DomainBase extends EppResource {
     }
 
     public B setRegistrant(Ref<ContactResource> registrant) {
-      getInstance().registrant = ReferenceUnion.create(registrant);
+      // Replace the registrant contact inside allContacts.
+      getInstance().allContacts = union(
+          getInstance().getContacts(),
+          DesignatedContact.create(Type.REGISTRANT, checkArgumentNotNull(registrant)));
       return thisCastToDerived();
     }
 
@@ -284,7 +310,13 @@ public abstract class DomainBase extends EppResource {
     }
 
     public B setContacts(ImmutableSet<DesignatedContact> contacts) {
-      getInstance().contacts = contacts;
+      checkArgument(all(contacts, not(IS_REGISTRANT)), "Registrant cannot be a contact");
+      // Replace the non-registrant contacts inside allContacts.
+      getInstance().allContacts = FluentIterable
+          .from(nullToEmpty(getInstance().allContacts))
+          .filter(IS_REGISTRANT)
+          .append(contacts)
+          .toSet();
       return thisCastToDerived();
     }
 

java/google/registry/model/domain/DomainCommand.java

@@ -510,20 +510,6 @@ public class DomainCommand {
           }}));
   }
 
-  /**
-   * EPP-inputable version of XML type for contact identifiers associated with a domain, which can
-   * be converted to a storable (and EPP-outputable) {@link DesignatedContact}.
-   *
-   * @see "http://tools.ietf.org/html/rfc5731#section-2.2"
-   */
-  static class ForeignKeyedDesignatedContact extends ImmutableObject {
-    @XmlAttribute(required = true)
-    DesignatedContact.Type type;
-
-    @XmlValue
-    String contactId;
-  }
-
   /** Exception to throw when referenced objects don't exist. */
   public static class InvalidReferencesException extends Exception {
     private final ImmutableSet<String> foreignKeys;

java/google/registry/model/domain/DomainResource.java

@@ -62,9 +62,9 @@ import javax.xml.bind.annotation.XmlType;
     "fullyQualifiedDomainName",
     "repoId",
     "status",
-    "registrant",
+    "marshalledRegistrant",
-    "contacts",
+    "marshalledContacts",
-    "nameservers",
+    "marshalledNameservers",
     "subordinateHosts",
     "currentSponsorClientId",
     "creationClientId",

java/google/registry/model/domain/ForeignKeyedDesignatedContact.java

@@ -0,0 +1,41 @@
+// Copyright 2016 The Domain Registry Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package google.registry.model.domain;
+
+import google.registry.model.ImmutableObject;
+
+import javax.xml.bind.annotation.XmlAttribute;
+import javax.xml.bind.annotation.XmlValue;
+
+/**
+ * EPP-compatible version of XML type for contact identifiers associated with a domain, which can
+ * be converted to a storable {@link DesignatedContact}.
+ *
+ * @see "http://tools.ietf.org/html/rfc5731#section-2.2"
+ */
+class ForeignKeyedDesignatedContact extends ImmutableObject {
+  @XmlAttribute(required = true)
+  DesignatedContact.Type type;
+
+  @XmlValue
+  String contactId;
+
+  static ForeignKeyedDesignatedContact create(DesignatedContact.Type type, String contactId) {
+    ForeignKeyedDesignatedContact instance = new ForeignKeyedDesignatedContact();
+    instance.type = type;
+    instance.contactId = contactId;
+    return instance;
+  }
+}

java/google/registry/model/domain/ReferenceUnion.java

@@ -20,17 +20,11 @@ import com.googlecode.objectify.annotation.Index;
 
 import google.registry.model.EppResource;
 import google.registry.model.ImmutableObject;
-import google.registry.model.contact.ContactResource;
-import google.registry.model.host.HostResource;
-
-import javax.xml.bind.annotation.adapters.XmlAdapter;
 
 /**
  * Legacy shell of a "union" type to represent referenced objects as either a foreign key or as a
  * link to another object in the datastore. In its current form it merely wraps a {@link Ref}.
  *
- * <p>This type always marshals as the "foreign key". We no longer use this type for unmarshalling.
- *
  * @param <T> the type being referenced
  */
 @Embed
@@ -43,27 +37,6 @@ public class ReferenceUnion<T extends EppResource> extends ImmutableObject {
     return linked;
   }
 
-  /** An adapter that marshals the linked {@link Ref} as its loaded foreign key. */
-  public static class Adapter<T extends EppResource>
-      extends XmlAdapter<String, ReferenceUnion<T>> {
-
-    @Override
-    public ReferenceUnion<T> unmarshal(String foreignKey) throws Exception {
-      throw new UnsupportedOperationException();
-    }
-
-    @Override
-    public String marshal(ReferenceUnion<T> reference) throws Exception {
-      return reference.getLinked().get().getForeignKey();
-    }
-  }
-
-  /** An adapter for references to contacts. */
-  static class ContactReferenceUnionAdapter extends Adapter<ContactResource>{}
-
-  /** An adapter for references to hosts. */
-  static class HostReferenceUnionAdapter extends Adapter<HostResource>{}
-
   public static <T extends EppResource> ReferenceUnion<T> create(Ref<T> linked) {
     ReferenceUnion<T> instance = new ReferenceUnion<>();
     instance.linked = linked;

java/google/registry/model/domain/package-info.java

@@ -19,13 +19,9 @@
 @XmlAccessorType(XmlAccessType.FIELD)
 @XmlJavaTypeAdapters({
     @XmlJavaTypeAdapter(UtcDateTimeAdapter.class),
-    @XmlJavaTypeAdapter(ContactReferenceUnionAdapter.class),
-    @XmlJavaTypeAdapter(HostReferenceUnionAdapter.class),
     @XmlJavaTypeAdapter(DateAdapter.class)})
 package google.registry.model.domain;
 
-import google.registry.model.domain.ReferenceUnion.ContactReferenceUnionAdapter;
-import google.registry.model.domain.ReferenceUnion.HostReferenceUnionAdapter;
 import google.registry.xml.DateAdapter;
 import google.registry.xml.UtcDateTimeAdapter;
 

javatests/google/registry/flows/domain/DomainUpdateFlowTest.java

@@ -34,6 +34,7 @@ import static google.registry.testing.HistoryEntrySubject.assertAboutHistoryEntr
 import static google.registry.testing.TaskQueueHelper.assertDnsTasksEnqueued;
 import static org.joda.money.CurrencyUnit.USD;
 
+import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.ImmutableSet;
 
@@ -380,20 +381,19 @@ public class DomainUpdateFlowTest extends ResourceFlowTestCase<DomainUpdateFlow,
         nameservers.add(Ref.create(host));
       }
     }
-    ImmutableSet.Builder<DesignatedContact> contacts = new ImmutableSet.Builder<>();
+    ImmutableList.Builder<DesignatedContact> contactsBuilder = new ImmutableList.Builder<>();
     for (int i = 0; i < 8; i++) {
-      ContactResource contact = persistActiveContact(String.format("max_test_%d", i));
+      contactsBuilder.add(
-      if (i < 4) {
+          DesignatedContact.create(
-        contacts.add(
+              DesignatedContact.Type.values()[i % 4],
-            DesignatedContact.create(
+              Ref.create(persistActiveContact(String.format("max_test_%d", i)))));
-                DesignatedContact.Type.values()[i],
-                Ref.create(contact)));
-      }
     }
+    ImmutableList<DesignatedContact> contacts = contactsBuilder.build();
     persistResource(
         reloadResourceByUniqueId().asBuilder()
             .setNameservers(nameservers.build())
-            .setContacts(contacts.build())
+            .setContacts(ImmutableSet.copyOf(contacts.subList(0, 3)))
+            .setRegistrant(contacts.get(3).getContactRef())
             .build());
     clock.advanceOneMilli();
     assertTransactionalFlow(true);