Fetch only the challenges needed in getAuthz. (#3597)

In `getAuthorizations`, we had a single loop to both select the freshest authz and
fetch challenges corresponding to authzs. This meant that in some cases, we
would fetch challenges only to throw them away. Since each challenge fetch is a
DB round trip, this would cause extreme slowness when called for domains that
have a large number of authorizations.

This change splits that into two loops: One to select the freshest authzs, and
another to fetch challenges for those authzs.

sa/sa.go

@@ -29,6 +29,7 @@ import (
 )
 
 type certCountFunc func(domain string, earliest, latest time.Time) (int, error)
+type getChallengesFunc func(authID string) ([]core.Challenge, error)
 
 // SQLStorageAuthority defines a Storage Authority
 type SQLStorageAuthority struct {
@@ -42,9 +43,10 @@ type SQLStorageAuthority struct {
 	// threads).
 	parallelismPerRPC int
 
-	// We use a function type here so we can mock out this internal function in
+	// We use function types here so we can mock out this internal function in
 	// unittests.
 	countCertificatesByName certCountFunc
+	getChallenges           getChallengesFunc
 }
 
 func digest256(data []byte) []byte {
@@ -106,6 +108,7 @@ func NewSQLStorageAuthority(
 	}
 
 	ssa.countCertificatesByName = ssa.countCertificatesByNameImpl
+	ssa.getChallenges = ssa.getChallengesImpl
 
 	return ssa, nil
 }
@@ -1947,16 +1950,17 @@ func (ssa *SQLStorageAuthority) getAuthorizations(
 		}
 		existing, present := byName[auth.Identifier.Value]
 		if !present || auth.Expires.After(*existing.Expires) {
-			// Retrieve challenges for the authz
-			auth.Challenges, err = ssa.getChallenges(auth.ID)
-			if err != nil {
-				return nil, err
-			}
-
 			byName[auth.Identifier.Value] = auth
 		}
 	}
 
+	for _, auth := range byName {
+		// Retrieve challenges for the authz
+		if auth.Challenges, err = ssa.getChallenges(auth.ID); err != nil {
+			return nil, err
+		}
+	}
+
 	return byName, nil
 }
 
@@ -2062,7 +2066,7 @@ func (ssa *SQLStorageAuthority) AddPendingAuthorizations(ctx context.Context, re
 	return &sapb.AuthorizationIDs{Ids: ids}, nil
 }
 
-func (ssa *SQLStorageAuthority) getChallenges(authID string) ([]core.Challenge, error) {
+func (ssa *SQLStorageAuthority) getChallengesImpl(authID string) ([]core.Challenge, error) {
 	var challObjs []challModel
 	_, err := ssa.dbMap.Select(
 		&challObjs,

sa/sa_test.go

@@ -2184,3 +2184,57 @@ func TestStatusForOrder(t *testing.T) {
 	}
 
 }
+
+// Check that getAuthorizations is fast enough; that is, it shouldn't retrieve
+// challenges for authorizations it won't return (which has been a cause of
+// slowness when there are many authorizations for the same domain name).
+func TestGetAuthorizationsFast(t *testing.T) {
+	sa, fc, cleanUp := initSA(t)
+	defer cleanUp()
+
+	ctx := context.Background()
+	reg := satest.CreateWorkingRegistration(t, sa)
+
+	expires := fc.Now().Add(time.Hour)
+
+	makeAuthz := func(s string) {
+		_, err := sa.NewPendingAuthorization(ctx, core.Authorization{
+			RegistrationID: reg.ID,
+			Expires:        &expires,
+			Status:         core.StatusPending,
+			Identifier: core.AcmeIdentifier{
+				Type:  core.IdentifierDNS,
+				Value: s,
+			},
+		})
+		test.AssertNotError(t, err, "making pending authz")
+	}
+
+	for i := 0; i < 10; i++ {
+		makeAuthz("example.com")
+		makeAuthz("www.example.com")
+		expires = expires.Add(time.Hour)
+	}
+
+	// Mock out getChallenges so we can count how many times it's called.
+	var challengeFetchCount int
+	sa.getChallenges = func(s string) ([]core.Challenge, error) {
+		challengeFetchCount++
+		return nil, nil
+	}
+
+	results, err := sa.getAuthorizations(ctx, pendingAuthorizationTable,
+		string(core.StatusPending), reg.ID, []string{"example.com", "www.example.com"},
+		fc.Now(), false)
+	test.AssertNotError(t, err, "getting authorizations")
+	if len(results) != 2 {
+		t.Fatalf("Wrong number of results. Expected 2, got %d", len(results))
+	}
+	if results["example.com"] == nil || results["www.example.com"] == nil {
+		t.Fatalf("Nil result for expected domain: %#v", results)
+	}
+	// We expect getChallenges to be called exactly once for each domain.
+	if challengeFetchCount != 2 {
+		t.Errorf("Wrong challenge fetch count: expected 2, got %d", challengeFetchCount)
+	}
+}