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

View File

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

View File

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

View File

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