Limit document count query runtime and fallback to estimation on expire.

Author: Jiri Malek

internal/repository/db/bridge.go

@@ -33,6 +33,11 @@ type MongoDbBridge struct {
 	initEpochs       *sync.Once
 }
 
+// docListCountAggregationTimeout represents a max duration of DB query executed to calculate
+// exact document count in filtered collection. If this duration is exceeded, the query fails
+// ad we fall back to full collection documents count estimation.
+const docListCountAggregationTimeout = 500 * time.Millisecond
+
 // intZero represents an empty big value.
 var intZero = new(big.Int)
 
@@ -216,3 +221,24 @@ func (db *MongoDbBridge) EstimateCount(col *mongo.Collection) (uint64, error) {
 	}
 	return uint64(val), nil
 }
+
+// listDocumentsCount tries to calculate precise documents count and if it's not counted in limited
+// time, use general estimation to speed up the loader.
+func (db *MongoDbBridge) listDocumentsCount(col *mongo.Collection, filter *bson.D) (int64, error) {
+	// try to count the proper way
+	total, err := col.CountDocuments(context.Background(), filter, options.Count().SetMaxTime(docListCountAggregationTimeout))
+	if err == nil {
+		return total, nil
+	}
+
+	// it failed in the limited time we gave it
+	db.log.Errorf("can not count documents properly; %s", err.Error())
+
+	// just estimate the whole collection size
+	total, err = col.EstimatedDocumentCount(context.Background())
+	if err != nil {
+		db.log.Errorf("can not count documents")
+		return 0, err
+	}
+	return total, nil
+}

internal/repository/db/transaction.go

@@ -178,7 +178,7 @@ func (db *MongoDbBridge) initTrxList(col *mongo.Collection, cursor *string, coun
 	}
 
 	// find how many transactions do we have in the database
-	total, err := col.CountDocuments(context.Background(), *filter)
+	total, err := db.listDocumentsCount(col, filter)
 	if err != nil {
 		db.log.Errorf("can not count transactions")
 		return nil, err