rtsp-video-analysis-system/ruoyi-video/src/main/java/com/ruoyi/video/service/MediaService.java

package com.ruoyi.video.service;

import cn.hutool.core.thread.ThreadUtil;
import cn.hutool.crypto.digest.MD5;
import com.ruoyi.video.common.ClientType;
import com.ruoyi.video.domain.dto.CameraDto;
import com.ruoyi.video.thread.MediaTransfer;
import com.ruoyi.video.thread.MediaTransferFlvByFFmpeg;
import com.ruoyi.video.thread.MediaTransferFlvByJavacv;
import io.netty.channel.ChannelHandlerContext;
import org.springframework.stereotype.Service;

import java.util.concurrent.ConcurrentHashMap;
import java.util.function.Consumer;

/**
 * 媒体服务，支持全局网络超时、读写超时、无人拉流持续时长自动关闭流等配置
 * @Author: orange
 * @CreateTime: 2025-01-16
 */
@Service
public class MediaService {

    /**
     * 缓存流转换线程
     */
    public static ConcurrentHashMap<String, MediaTransfer> cameras = new ConcurrentHashMap<>();

    /**
     * http-flv播放
     * @param cameraDto 摄像头配置
     * @param ctx       Netty上下文
     */
    public void playForHttp(CameraDto cameraDto, ChannelHandlerContext ctx) {

        if (cameras.containsKey(cameraDto.getMediaKey())) {
            MediaTransfer mediaConvert = cameras.get(cameraDto.getMediaKey());
            if (mediaConvert instanceof MediaTransferFlvByJavacv) {
                MediaTransferFlvByJavacv mediaTransferFlvByJavacv = (MediaTransferFlvByJavacv) mediaConvert;
                //如果当前已经用ffmpeg，则重新拉流
                if (cameraDto.isEnabledFFmpeg()) {
                    mediaTransferFlvByJavacv.setRunning(false);
                    cameras.remove(cameraDto.getMediaKey());
                    this.playForHttp(cameraDto, ctx);
                } else {
                    mediaTransferFlvByJavacv.addClient(ctx, ClientType.HTTP);
                }
            } else if (mediaConvert instanceof MediaTransferFlvByFFmpeg) {
                MediaTransferFlvByFFmpeg mediaTransferFlvByFFmpeg = (MediaTransferFlvByFFmpeg) mediaConvert;
                //如果当前已经用javacv，则关闭再重新拉流
                if (!cameraDto.isEnabledFFmpeg()) {
                    mediaTransferFlvByFFmpeg.stopFFmpeg();
                    cameras.remove(cameraDto.getMediaKey());
                    this.playForHttp(cameraDto, ctx);
                } else {
                    mediaTransferFlvByFFmpeg.addClient(ctx, ClientType.HTTP);
                }
            }

        } else {
            if (cameraDto.isEnabledFFmpeg()) {
                MediaTransferFlvByFFmpeg mediaft = new MediaTransferFlvByFFmpeg(cameraDto);
                mediaft.execute();
                cameras.put(cameraDto.getMediaKey(), mediaft);
                mediaft.addClient(ctx, ClientType.HTTP);
            } else {
                MediaTransferFlvByJavacv mediaConvert = new MediaTransferFlvByJavacv(cameraDto);
                cameras.put(cameraDto.getMediaKey(), mediaConvert);
                ThreadUtil.execute(mediaConvert);
                mediaConvert.addClient(ctx, ClientType.HTTP);
            }

        }
    }

    /**
     * ws-flv播放
     * @param cameraDto 摄像头配置
     * @param ctx       Netty上下文
     */
    public void playForWs(CameraDto cameraDto, ChannelHandlerContext ctx) {

        if (cameras.containsKey(cameraDto.getMediaKey())) {
            MediaTransfer mediaConvert = cameras.get(cameraDto.getMediaKey());
            if (mediaConvert instanceof MediaTransferFlvByJavacv) {
                MediaTransferFlvByJavacv mediaTransferFlvByJavacv = (MediaTransferFlvByJavacv) mediaConvert;
                //如果当前已经用ffmpeg，则重新拉流
                if (cameraDto.isEnabledFFmpeg()) {
                    mediaTransferFlvByJavacv.setRunning(false);
                    cameras.remove(cameraDto.getMediaKey());
                    this.playForWs(cameraDto, ctx);
                } else {
                    mediaTransferFlvByJavacv.addClient(ctx, ClientType.WEBSOCKET);
                }
            } else if (mediaConvert instanceof MediaTransferFlvByFFmpeg) {
                MediaTransferFlvByFFmpeg mediaTransferFlvByFFmpeg = (MediaTransferFlvByFFmpeg) mediaConvert;
                //如果当前已经用javacv，则关闭再重新拉流
                if (!cameraDto.isEnabledFFmpeg()) {
                    mediaTransferFlvByFFmpeg.stopFFmpeg();
                    cameras.remove(cameraDto.getMediaKey());
                    this.playForWs(cameraDto, ctx);
                } else {
                    mediaTransferFlvByFFmpeg.addClient(ctx, ClientType.WEBSOCKET);
                }
            }
        } else {
            if (cameraDto.isEnabledFFmpeg()) {
                MediaTransferFlvByFFmpeg mediaft = new MediaTransferFlvByFFmpeg(cameraDto);
                mediaft.execute();
                cameras.put(cameraDto.getMediaKey(), mediaft);
                mediaft.addClient(ctx, ClientType.WEBSOCKET);
            } else {
                MediaTransferFlvByJavacv mediaConvert = new MediaTransferFlvByJavacv(cameraDto);
                cameras.put(cameraDto.getMediaKey(), mediaConvert);
                ThreadUtil.execute(mediaConvert);
                mediaConvert.addClient(ctx, ClientType.WEBSOCKET);
            }
        }
    }

    /**
     * api播放
     * @param cameraDto 摄像头配置
     * @return 是否启动成功
     */
    public boolean playForApi(CameraDto cameraDto) {
        // 区分不同媒体
        String mediaKey = MD5.create().digestHex(cameraDto.getUrl());
        cameraDto.setMediaKey(mediaKey);
        cameraDto.setEnabledFlv(true);

        MediaTransfer mediaTransfer = cameras.get(cameraDto.getMediaKey());
        if (null == mediaTransfer) {
            if (cameraDto.isEnabledFFmpeg()) {
                MediaTransferFlvByFFmpeg mediaft = new MediaTransferFlvByFFmpeg(cameraDto);
                mediaft.execute();
                cameras.put(cameraDto.getMediaKey(), mediaft);
            } else {
                MediaTransferFlvByJavacv mediaConvert = new MediaTransferFlvByJavacv(cameraDto);
                cameras.put(cameraDto.getMediaKey(), mediaConvert);
                ThreadUtil.execute(mediaConvert);
            }
        }

        mediaTransfer = cameras.get(cameraDto.getMediaKey());
        //同步等待
        if (mediaTransfer instanceof MediaTransferFlvByJavacv) {
            MediaTransferFlvByJavacv mediaft = (MediaTransferFlvByJavacv) mediaTransfer;
            // 30秒还没true认为启动不了
            for (int i = 0; i < 60; i++) {
                if (mediaft.isRunning() && mediaft.isGrabberStatus() && mediaft.isRecorderStatus()) {
                    return true;
                }
                try {
                    Thread.sleep(500);
                } catch (InterruptedException e) {
                    // ignore
                }
            }
        } else if (mediaTransfer instanceof MediaTransferFlvByFFmpeg) {
            MediaTransferFlvByFFmpeg mediaft = (MediaTransferFlvByFFmpeg) mediaTransfer;
            // 30秒还没true认为启动不了
            for (int i = 0; i < 60; i++) {
                if (mediaft.isRunning()) {
                    return true;
                }
                try {
                    Thread.sleep(500);
                } catch (InterruptedException e) {
                    // ignore
                }
            }
        }
        return false;
    }

    /**
     * 关闭流
     * @param cameraDto 摄像头配置
     */
    public void closeForApi(CameraDto cameraDto) {
        cameraDto.setEnabledFlv(false);

        if (cameras.containsKey(cameraDto.getMediaKey())) {
            MediaTransfer mediaConvert = cameras.get(cameraDto.getMediaKey());
            if (mediaConvert instanceof MediaTransferFlvByJavacv) {
                MediaTransferFlvByJavacv mediaTransferFlvByJavacv = (MediaTransferFlvByJavacv) mediaConvert;
                mediaTransferFlvByJavacv.setRunning(false);
                cameras.remove(cameraDto.getMediaKey());
            } else if (mediaConvert instanceof MediaTransferFlvByFFmpeg) {
                MediaTransferFlvByFFmpeg mediaTransferFlvByFFmpeg = (MediaTransferFlvByFFmpeg) mediaConvert;
                mediaTransferFlvByFFmpeg.stopFFmpeg();
                cameras.remove(cameraDto.getMediaKey());
            }
        }
    }

    /* =========================== 新增便捷方法 =========================== */

    /** 直接从缓存取 MediaTransfer（可能是 FFmpeg 或 JavaCV）。不存在返回 null。 */
    public MediaTransfer getMedia(String mediaKey) {
        return cameras.get(mediaKey);
    }

    /** 只取 JavaCV 实例；如果不是 JavaCV 或不存在则返回 null。 */
    public MediaTransferFlvByJavacv getJavacv(String mediaKey) {
        MediaTransfer mt = cameras.get(mediaKey);
        return (mt instanceof MediaTransferFlvByJavacv) ? (MediaTransferFlvByJavacv) mt : null;
    }

    /**
     * 取或启动 JavaCV 实例：
     * - 已有 JavaCV：直接返回
     * - 已有 FFmpeg：先停止 FFmpeg，再切换 JavaCV
     * - 不存在：启动 JavaCV
     *
     * @param cameraDto   需包含 url / mediaKey（mediaKey 为空则用 url 的 MD5 生成）
     * @param beforeStart 启动前对 cameraDto 做一次定制（可 null），例如 dto -> dto.setEnableDetection(true)
     */
    public MediaTransferFlvByJavacv getOrStartJavacv(CameraDto cameraDto, Consumer<CameraDto> beforeStart) {
        // 兜底 mediaKey
        if (cameraDto.getMediaKey() == null || cameraDto.getMediaKey().isEmpty()) {
            String mediaKey = MD5.create().digestHex(cameraDto.getUrl());
            cameraDto.setMediaKey(mediaKey);
        }

        MediaTransfer mt = cameras.get(cameraDto.getMediaKey());
        if (mt instanceof MediaTransferFlvByJavacv) {
            return (MediaTransferFlvByJavacv) mt;
        }

        // 若已存在 FFmpeg 实例，先停掉
        if (mt instanceof MediaTransferFlvByFFmpeg) {
            ((MediaTransferFlvByFFmpeg) mt).stopFFmpeg();
            cameras.remove(cameraDto.getMediaKey());
        }

        // 启动 JavaCV
        if (beforeStart != null) beforeStart.accept(cameraDto);
        MediaTransferFlvByJavacv mediaConvert = new MediaTransferFlvByJavacv(cameraDto);
        cameras.put(cameraDto.getMediaKey(), mediaConvert);
        ThreadUtil.execute(mediaConvert);
        return mediaConvert;
    }

    /** 可选：根据 mediaKey 强制停止并移除（两种实现都兼容） */
    public void stopByMediaKey(String mediaKey) {
        MediaTransfer mt = cameras.get(mediaKey);
        if (mt instanceof MediaTransferFlvByJavacv) {
            ((MediaTransferFlvByJavacv) mt).setRunning(false);
        } else if (mt instanceof MediaTransferFlvByFFmpeg) {
            ((MediaTransferFlvByFFmpeg) mt).stopFFmpeg();
        }
        cameras.remove(mediaKey);
    }
}