UpdateInstalledExtensionCountScheduler

Update the counter of installed extensions

import groovy.json.JsonOutput
import groovy.json.JsonSlurper

import org.xwiki.activeinstalls2.ActiveInstallsConfiguration

// Instances running XWiki < 14.5 send their pings to the Active Installs 1 Elasticsearch instance, newer ones send
// them to the Active Installs 2 instance. Both have to be queried and their results added up, otherwise only a
// fraction of the running instances is counted.

/**
 * @return the number of distinct instances having pinged Active Installs 2 during the last day, keyed by extension id
 */
def getActiveInstalls2Counts()
{
  def query = [
    'size': 0,
    'query': [
      'range': [
        'date.current': ['gte': 'now-1d']
      ]
    ],
    'aggs': [
      'extensions': [
        // "extensions" is mapped as a nested type so that each entry's id/version/features stay indexed together,
        // which means it can only be reached through a nested aggregation.
        'nested': ['path': 'extensions'],
        'aggs': [
          'ids': [
            'terms': ['field': 'extensions.id', 'size': 10000],
            'aggs': [
              // The instance id is on the root document and not on the nested extension, so step back out to count it.
              'instances': [
                'reverse_nested': [:],
                'aggs': [
                  'count': ['cardinality': ['field': 'distribution.instanceId']]
                ]
              ]
            ]
          ]
        ]
      ]
    ]
  ]

  // The Active Installs 2 script service exposes no aggregation API, so query the Elasticsearch REST endpoint
  // directly. A single request returns the count for every extension, which is much cheaper than the one request per
  // extension that Active Installs 1 requires. Once this wiki runs XWiki 18.7 or later, this whole method can be
  // replaced by a call to services.activeinstalls2.countDistinctInstallsByExtension().
  def configuration = services.component.getInstance(ActiveInstallsConfiguration.class)
  def connection = new URL("${configuration.pingInstanceURL}/_search").openConnection()
  connection.requestMethod = 'POST'
  connection.doOutput = true
  connection.connectTimeout = 30000
  connection.readTimeout = 300000
  connection.setRequestProperty('Content-Type', 'application/json')
  // The default JDK user agent gets served a Cloudflare challenge instead of the query result.
  connection.setRequestProperty('User-Agent', configuration.userAgent)
  connection.outputStream.withWriter('UTF-8') { writer -> writer << JsonOutput.toJson(query) }

  if (connection.responseCode != 200) {
    throw new Exception("Active Installs 2 query failed with [${connection.responseCode}]: "
      + "[${connection.errorStream?.getText('UTF-8')}]")
  }

  def response
  connection.inputStream.withReader('UTF-8') { reader -> response = new JsonSlurper().parse(reader) }

  def buckets = response.aggregations.extensions.ids.buckets
  if (response.aggregations.extensions.ids.sum_other_doc_count > 0) {
    // A terms aggregation silently drops the extensions that don't fit in its size, which would under-report them.
    throw new Exception("Too many extension ids returned by Active Installs 2, raise the aggregation size above "
      + "[${buckets.size()}]")
  }

  def counts = [:]
  buckets.each { bucket -> counts.put(bucket.key, bucket.instances.count.value as long) }
  return counts
}

/**
 * @return the number of distinct instances having pinged Active Installs 1 during the last day for the passed
 *         extension. Contrary to Active Installs 2, the whole set cannot be aggregated in one query: the old
 *         Elasticsearch instance holds too much data for a terms aggregation on the extension ids and breaks its
 *         circuit breaker, so one query per extension is needed.
 */
def getActiveInstalls1Count(extensionId)
{
  def query = [
    'query': [
      'bool': [
        'must': [
          ['range': ['_timestamp': ['gte': 'now-1d']]],
          // The extension id is passed as a term value instead of being interpolated into a query string, so that an
          // id containing Lucene syntax cannot change what the query matches.
          ['term': ['extensions.id': extensionId]]
        ]
      ]
    ],
    'aggs': [
      'instanceId_count': [
        'cardinality': [
          'field': 'instanceId'
        ]
      ]
    ]
  ]

  def result = services.activeinstalls.searchInstalls('install2', JsonOutput.toJson(query), ['search_type' : 'count'])
  if (result == null) {
    throw services.activeinstalls.lastError ?:
      new Exception("Active Installs 1 returned no result for extension [$extensionId]")
  }

  return result.getAsJsonObject('aggregations').getAsJsonObject('instanceId_count').getAsJsonPrimitive('value')
    .getAsLong()
}

def updateCount(extensionId, documentName, activeInstalls2Counts)
{
  def pxwiki = xwiki.getXWiki()
  def pxcontext = xcontext
  def document = pxwiki.getDocument(documentName, pxcontext)
  def extensionObject = document.getObject('ExtensionCode.ExtensionClass')

  // Extensions not displaying their install count don't need to be counted at all.
  if (extensionObject.getIntValue('showInstalledCount') == 0) {
    return
  }

  def count = (activeInstalls2Counts.get(extensionId) ?: 0L) + getActiveInstalls1Count(extensionId)

  if (extensionObject.getLongValue('installedCount') != count) {
    extensionObject.set('installedCount', count, pxcontext)

    if (document.comment == 'Update install count') {
      // Don't increment version (otherwise we could end up with a lot of those)
      document.setMetaDataDirty(false)
      document.setContentDirty(false)
    }

    // Save
    pxwiki.saveDocument(document, 'Update install count', true, pxcontext)
  }
}

def logger = org.slf4j.LoggerFactory.getLogger('ExtensionCode.UpdateInstalledExtensionCountScheduler')

logger.info('Starting updating installed extensions count')

// Fail the whole job rather than saving Active Installs 1 counts only, which would look like a massive drop.
def activeInstalls2Counts = getActiveInstalls2Counts()
logger.info('Retrieved Active Installs 2 counts for [{}] extensions', activeInstalls2Counts.size())

def extensions = services.query.xwql("select doc.fullName, extension.id from Document doc, doc.object(ExtensionCode.ExtensionClass) as extension").execute()
for (extension in extensions)
{
  logger.info('Updating installed extension [{}] count', extension[0])

  try {
    updateCount(extension[1], extension[0], activeInstalls2Counts)
  } catch (Exception e) {
    logger.error('Failed to update installed count for extension [{}]', extension[0], e)
  }
}

logger.info('Installed extensions count update is finished')