feat: cache media uploads

This commit is contained in:
Adhiraj Singh
2021-11-10 19:22:00 +05:30
parent cfa7e8ec66
commit 3e54741042
4 changed files with 34 additions and 14 deletions

View File

@@ -379,7 +379,7 @@ export const makeMessagesSocket = (config: SocketConfig) => {
return msgId
}
const waUploadToServer: WAMediaUploadFunction = async(stream, { mediaType, fileEncSha256B64 }) => {
const waUploadToServer: WAMediaUploadFunction = async(stream, { mediaType, fileEncSha256B64, timeoutMs }) => {
// send a query JSON to obtain the url & auth token to upload our media
let uploadInfo = await refreshMediaConn(false)
@@ -399,7 +399,8 @@ export const makeMessagesSocket = (config: SocketConfig) => {
agent: {
https: config.agent
},
body: stream
body: stream,
timeout: timeoutMs
}
)
const result = JSON.parse(responseText)
@@ -456,10 +457,11 @@ export const makeMessagesSocket = (config: SocketConfig) => {
{
...options,
logger,
userJid: userJid,
userJid,
// multi-device does not have this yet
//getUrlInfo: generateUrlInfo,
upload: waUploadToServer
upload: waUploadToServer,
mediaCache: config.mediaCache,
}
)
const isDeleteMsg = 'delete' in content && !!content.delete

View File

@@ -1,6 +1,7 @@
import type { ReadStream } from "fs"
import type { Logger } from "pino"
import type { URL } from "url"
import type NodeCache from "node-cache"
import type { GroupMetadata } from "./GroupMetadata"
import { proto } from '../../WAProto'
@@ -121,18 +122,22 @@ export type MiscMessageGenerationOptions = {
quoted?: WAMessage
/** disappearing messages settings */
ephemeralExpiration?: number | string
mediaUploadTimeoutMs?: number
}
export type MessageGenerationOptionsFromContent = MiscMessageGenerationOptions & {
userJid: string
}
export type WAMediaUploadFunction = (readStream: ReadStream, opts: { fileEncSha256B64: string, mediaType: MediaType }) => Promise<{ mediaUrl: string }>
export type WAMediaUploadFunction = (readStream: ReadStream, opts: { fileEncSha256B64: string, mediaType: MediaType, timeoutMs?: number }) => Promise<{ mediaUrl: string }>
export type MediaGenerationOptions = {
logger?: Logger
upload: WAMediaUploadFunction
/** cache media so it does not have to be uploaded again */
mediaCache?: (url: string) => Promise<WAGenericMediaMessage> | WAGenericMediaMessage
mediaCache?: NodeCache
mediaUploadTimeoutMs?: number
}
export type MessageContentGenerationOptions = MediaGenerationOptions & {
getUrlInfo?: (text: string) => Promise<WAUrlInfo>

View File

@@ -48,6 +48,8 @@ export type SocketConfig = {
emitOwnEvents: boolean
/** provide a cache to store a user's device list */
userDevicesCache?: NodeCache
/** provide a cache to store media, so does not have to be re-uploaded */
mediaCache?: NodeCache
}
export enum DisconnectReason {

View File

@@ -67,13 +67,19 @@ export const prepareWAMessageMedia = async(
[mediaType]: undefined,
media: message[mediaType]
}
// check if cacheable + generate cache key
const cacheableKey = typeof uploadData.media === 'object' &&
('url' in uploadData.media) &&
!!uploadData.media.url &&
!!options.mediaCache && (
// generate the key
mediaType + ':' + uploadData.media.url!.toString()
)
// check for cache hit
if(typeof uploadData.media === 'object' && 'url' in uploadData.media) {
const result = !!options.mediaCache && await options.mediaCache!(uploadData.media.url?.toString())
if(result) {
return WAProto.Message.fromObject({
[`${mediaType}Message`]: result
})
if(cacheableKey) {
const mediaBuff: Buffer = options.mediaCache!.get(cacheableKey)
if(mediaBuff) {
return WAProto.Message.decode(mediaBuff)
}
}
if(mediaType === 'document' && !uploadData.fileName) {
@@ -114,7 +120,7 @@ export const prepareWAMessageMedia = async(
}
const {mediaUrl} = await options.upload(
createReadStream(encBodyPath),
{ fileEncSha256B64, mediaType }
{ fileEncSha256B64, mediaType, timeoutMs: options.mediaUploadTimeoutMs }
)
// remove tmp files
await Promise.all(
@@ -138,7 +144,12 @@ export const prepareWAMessageMedia = async(
}
)
}
return WAProto.Message.fromObject(content)
const obj = WAProto.Message.fromObject(content)
if(cacheableKey) {
options.mediaCache!.set(cacheableKey, WAProto.Message.encode(obj))
}
return obj
}
export const prepareDisappearingMessageSettingContent = (ephemeralExpiration?: number) => {
ephemeralExpiration = ephemeralExpiration || 0