136 lines
4.7 KiB
Kotlin
136 lines
4.7 KiB
Kotlin
package io.legato.kazusa.help
|
|
|
|
import io.legato.kazusa.data.entities.BaseSource
|
|
import io.legato.kazusa.exception.ConcurrentException
|
|
import io.legato.kazusa.model.analyzeRule.AnalyzeUrl.ConcurrentRecord
|
|
import kotlinx.coroutines.delay
|
|
|
|
class ConcurrentRateLimiter(val source: BaseSource?) {
|
|
|
|
companion object {
|
|
private val concurrentRecordMap = hashMapOf<String, ConcurrentRecord>()
|
|
}
|
|
|
|
/**
|
|
* 开始访问,并发判断
|
|
*/
|
|
@Throws(ConcurrentException::class)
|
|
private fun fetchStart(): ConcurrentRecord? {
|
|
source ?: return null
|
|
val concurrentRate = source.concurrentRate
|
|
if (concurrentRate.isNullOrEmpty() || concurrentRate == "0") {
|
|
return null
|
|
}
|
|
val rateIndex = concurrentRate.indexOf("/")
|
|
var fetchRecord = concurrentRecordMap[source.getKey()]
|
|
if (fetchRecord == null) {
|
|
synchronized(concurrentRecordMap) {
|
|
fetchRecord = concurrentRecordMap[source.getKey()]
|
|
if (fetchRecord == null) {
|
|
fetchRecord = ConcurrentRecord(rateIndex > 0, System.currentTimeMillis(), 1)
|
|
concurrentRecordMap[source.getKey()] = fetchRecord
|
|
return fetchRecord
|
|
}
|
|
}
|
|
}
|
|
val waitTime: Int = synchronized(fetchRecord!!) {
|
|
try {
|
|
if (!fetchRecord.isConcurrent) {
|
|
//并发控制非 次数/毫秒
|
|
if (fetchRecord.frequency > 0) {
|
|
//已经有访问线程,直接等待
|
|
return@synchronized concurrentRate.toInt()
|
|
}
|
|
//没有线程访问,判断还剩多少时间可以访问
|
|
val nextTime = fetchRecord.time + concurrentRate.toInt()
|
|
if (System.currentTimeMillis() >= nextTime) {
|
|
fetchRecord.time = System.currentTimeMillis()
|
|
fetchRecord.frequency = 1
|
|
return@synchronized 0
|
|
}
|
|
return@synchronized (nextTime - System.currentTimeMillis()).toInt()
|
|
} else {
|
|
//并发控制为 次数/毫秒
|
|
val sj = concurrentRate.substring(rateIndex + 1)
|
|
val nextTime = fetchRecord.time + sj.toInt()
|
|
if (System.currentTimeMillis() >= nextTime) {
|
|
//已经过了限制时间,重置开始时间
|
|
fetchRecord.time = System.currentTimeMillis()
|
|
fetchRecord.frequency = 1
|
|
return@synchronized 0
|
|
}
|
|
val cs = concurrentRate.substring(0, rateIndex)
|
|
if (fetchRecord.frequency > cs.toInt()) {
|
|
return@synchronized (nextTime - System.currentTimeMillis()).toInt()
|
|
} else {
|
|
fetchRecord.frequency += 1
|
|
return@synchronized 0
|
|
}
|
|
}
|
|
} catch (_: Exception) {
|
|
return@synchronized 0
|
|
}
|
|
}
|
|
if (waitTime > 0) {
|
|
throw ConcurrentException(
|
|
"根据并发率还需等待${waitTime}毫秒才可以访问",
|
|
waitTime = waitTime
|
|
)
|
|
}
|
|
return fetchRecord
|
|
}
|
|
|
|
/**
|
|
* 访问结束
|
|
*/
|
|
fun fetchEnd(concurrentRecord: ConcurrentRecord?) {
|
|
if (concurrentRecord != null && !concurrentRecord.isConcurrent) {
|
|
synchronized(concurrentRecord) {
|
|
concurrentRecord.frequency -= 1
|
|
}
|
|
}
|
|
}
|
|
|
|
/**
|
|
* 获取并发记录，若处于并发限制状态下则会等待
|
|
*/
|
|
suspend fun getConcurrentRecord(): ConcurrentRecord? {
|
|
while (true) {
|
|
try {
|
|
return fetchStart()
|
|
} catch (e: ConcurrentException) {
|
|
delay(e.waitTime.toLong())
|
|
}
|
|
}
|
|
}
|
|
|
|
fun getConcurrentRecordBlocking(): ConcurrentRecord? {
|
|
while (true) {
|
|
try {
|
|
return fetchStart()
|
|
} catch (e: ConcurrentException) {
|
|
Thread.sleep(e.waitTime.toLong())
|
|
}
|
|
}
|
|
}
|
|
|
|
suspend inline fun <T> withLimit(block: () -> T): T {
|
|
val concurrentRecord = getConcurrentRecord()
|
|
try {
|
|
return block()
|
|
} finally {
|
|
fetchEnd(concurrentRecord)
|
|
}
|
|
}
|
|
|
|
inline fun <T> withLimitBlocking(block: () -> T): T {
|
|
val concurrentRecord = getConcurrentRecordBlocking()
|
|
try {
|
|
return block()
|
|
} finally {
|
|
fetchEnd(concurrentRecord)
|
|
}
|
|
}
|
|
|
|
}
|