/**
 * Copyright 2018 Twitter, Inc.
 * Licensed under the Apache License, Version 2.0
 * http://www.apache.org/licenses/LICENSE-2.0
 */ import { Context , SqrlKey } from "sqrl" ; import { getBucketSize , getBucketTimeForTimeMs , getBucketKey , getCurrentBucketExpirySeconds , getAllBucketKeys , } from "./BucketedKeys" ; import { TOTAL_COUNT_EXPIRY_SEC } from "./RedisCountService" ; import { CountUniqueService } from "../Services" ; import { RedisInterface , createRedisKey } from "./RedisInterface" ; const NUM_BUCKETS = 10 ; export class RedisApproxCountUniqueService implements CountUniqueService { constructor ( private redis , private prefix ) { } private async bumpTotal ( ctx , key , sortedHashes ) { const redisKey = createRedisKey ( ctx . requireDatabaseSet ( ) , this . prefix , key . getHex ( ) ) ; await Promise . all ( [ this . redis . pfadd ( ctx , redisKey , sortedHashes ) , this . redis . expire ( ctx , redisKey , TOTAL_COUNT_EXPIRY_SEC ) , ] ) ; } async bump ( ctx , props ) { const { at , key , sortedHashes , windowMs } = props ; if ( windowMs === null ) { return this . bumpTotal ( ctx , key , sortedHashes ) ; } const bucketSize = getBucketSize ( windowMs , NUM_BUCKETS ) ; const currentBucket = getBucketTimeForTimeMs ( at , bucketSize ) ; const redisKey = getBucketKey ( ctx . requireDatabaseSet ( ) , this . prefix , key . getHex ( ) , windowMs , currentBucket ) ; await this . redis . pfadd ( ctx , redisKey , sortedHashes ) ; await this . redis . expire ( ctx , redisKey , getCurrentBucketExpirySeconds ( windowMs , bucketSize ) ) ; } async fetchHashes ( ctx , props ) { throw new Error ( "fetchHashes() is not implemented" ) ; } async fetchTotal ( ctx , key , at , windowMs , extraHashesKey ) { const redisKey = createRedisKey ( ctx . requireDatabaseSet ( ) , this . prefix , key . getHex ( ) ) ; return this . redis . pfcount ( ctx , extraHashesKey ? [ redisKey , extraHashesKey ] : [ redisKey ] ) ; } async fetchCounts ( ctx , props ) { const { at , keys , windowMs , addHashes } = props ; const databaseSet = ctx . requireDatabaseSet ( ) ; let extraHashesKey = null ; if ( addHashes . length ) { extraHashesKey = createRedisKey ( databaseSet , this . prefix , "temp" , ... addHashes ) ; Promise . all ( [ this . redis . pfadd ( ctx , extraHashesKey , addHashes ) , this . redis . expire ( ctx , extraHashesKey , 1 ) , ] ) . catch ( ( err ) => { ctx . warn ( { } , "Error on countUnique() temporary hash insertion: " + err . toString ( ) ) ; } ) ; } return Promise . all ( keys . map ( async ( key ) => { if ( windowMs === null ) { return this . fetchTotal ( ctx , key , at , windowMs , extraHashesKey ) ; } const keys = getAllBucketKeys ( databaseSet , this . prefix , key , at , windowMs , NUM_BUCKETS ) ; if ( extraHashesKey ) { keys . push ( extraHashesKey ) ; } return this . redis . pfcount ( ctx , keys ) ; } ) ) ; } }