Picasso downloader that uses DiskLruCache and last Uri last path segment as key (useful for pre-signed url that changes frequently)

DiskLruCachePicassoDownloader.kt

import android.net.Uri
import app.choco.tracking.logger.Logger
import com.jakewharton.disklrucache.DiskLruCache
import com.squareup.picasso.Downloader
import java.io.File
import java.io.IOException
import java.io.InputStream
import java.net.HttpURLConnection
import java.net.URL
import java.util.Locale
import okhttp3.MediaType
import okhttp3.Protocol
import okhttp3.Request
import okhttp3.Response
import okhttp3.ResponseBody
import okio.BufferedSource
import okio.buffer
import okio.source

private const val DISK_CACHE_INDEX = 0
private const val KEY_LIMIT = 36 // 36 is the length of Version-4 UUIDs

class DiskLruCachePicassoDownloader(
    private val parentDirectory: File,
    private val cacheName: String,
    private val cacheMaxSize: Int,
    private val appVersion: Int = 1
) : Downloader {
    private val diskCacheLock = Object()
    private var diskCacheStarting = true
    private lateinit var _diskCache: DiskLruCache

    private val diskCache: DiskLruCache
        get() {
            initDiskCache()
            synchronized(diskCacheLock) {
                while (diskCacheStarting) {
                    try {
                        diskCacheLock.wait()
                    } catch (e: InterruptedException) {
                    }
                }
                return _diskCache
            }
        }

    init {
        initDiskCache()
    }

    private fun initDiskCache() {
        synchronized(diskCacheLock) {
            if (this::_diskCache.isInitialized && !_diskCache.isClosed) return
            try {
                val diskCacheDir = getCacheDirectory()
                if (!diskCacheDir.exists()) {
                    diskCacheDir.mkdirs()
                }
                _diskCache = DiskLruCache.open(
                    diskCacheDir,
                    appVersion,
                    1,
                    cacheMaxSize.toLong()
                )
            } catch (e: IOException) {
                Logger.e(e, "initDiskCache")
            } finally {
                diskCacheStarting = false
                diskCacheLock.notifyAll()
            }
        }
    }

    override fun load(request: Request): Response {
        val url = request.url().toString()
        val cacheControl = request.cacheControl()

        if (cacheControl.onlyIfCached() || !cacheControl.noCache()) {
            getFromDiskCache(url)?.let {
                return getResponse(request, it)
            }
        }

        val inputStream = downloadData(url)
        if (!request.cacheControl().noStore()) {
            writeToDiskCache(url, inputStream)
        }
        return getResponse(request, inputStream)
    }

    override fun shutdown() {
        try {
            diskCache.close()
        } catch (ignored: IOException) {
        }
    }

    private fun downloadData(url: String?): InputStream {
        val connection = URL(url).openConnection() as HttpURLConnection
        connection.useCaches = true
        val responseCode = connection.responseCode
        if (responseCode >= 300) {
            connection.disconnect()
            throw IOException("$responseCode ${connection.responseMessage}")
        }
        return connection.inputStream
    }

    private fun writeToDiskCache(url: String, inputStream: InputStream) {
        try {
            val key = hashKeyForDisk(url)
            val editor = diskCache.edit(key) ?: return
            editor.newOutputStream(DISK_CACHE_INDEX).use { outputStream ->
                inputStream.use {
                    it.copyTo(outputStream)
                }
            }
            editor.commit()
        } catch (e: Throwable) {
            Logger.e(e, "writeToDiskCache")
        }
    }

    private fun getFromDiskCache(url: String): InputStream? {
        return try {
            val key = hashKeyForDisk(url)
            val snapshot = diskCache[key] ?: return null
            snapshot.getInputStream(DISK_CACHE_INDEX)
        } catch (e: Throwable) {
            Logger.e(e, "loadFromDiskCache")
            null
        }
    }

    private fun getResponse(request: Request, inputStream: InputStream): Response {
        return Response.Builder()
            .body(
                InputStreamResponseBody(
                    MediaType.get("image/jpeg"),
                    inputStream
                )
            )
            .request(request)
            .code(200)
            .protocol(Protocol.HTTP_1_1)
            .message("Cached")
            .build()
    }

    /**
     * Key returned must match the regex [a-z0-9_-] and max of 64 characters in length
     * @see [DiskLruCache.LEGAL_KEY_PATTERN]
     */
    private fun hashKeyForDisk(url: String): String {
        return Uri.parse(url).lastPathSegment?.let {
            if (it.matches("[a-z0-9_-]".toRegex())) {
                it
            } else {
                it.toLowerCase(Locale.ENGLISH)
                    .replace("[^a-z0-9_-]".toRegex(), "")
                    .take(KEY_LIMIT)
            }
        } ?: url.hashCode().toString().take(KEY_LIMIT)
    }

    private fun getCacheDirectory(): File {
        return File(parentDirectory, cacheName)
    }

    /**
     * Clears the created cache directory included all the created files
     */
    fun clearCache(): Boolean {
        return getCacheDirectory().deleteRecursively()
    }
}

class InputStreamResponseBody(
    private val contentType: MediaType,
    private val inputStream: InputStream
) : ResponseBody() {

    override fun contentType(): MediaType? {
        return contentType
    }

    override fun contentLength(): Long {
        return -1
    }

    override fun source(): BufferedSource {
        return inputStream.source().buffer()
    }
}