MediaSource/MediaSourceRtsp.cpp

#include <thread>

#include <windows.h>

#include "MediaSourceRtsp.h"

#pragma comment(lib, "avformat.lib")
#pragma comment(lib, "avcodec.lib")
#pragma comment(lib, "avdevice.lib")
#pragma comment(lib, "avutil.lib")

CMediaSourceRtsp::CMediaSourceRtsp()
{
	outputContext = nullptr;
	inputContext = nullptr;
	m_nStartTime = 0;
	m_nWriteflag = false;
	m_timeout = 2;
	m_NetCount = 0;
	m_nVideoStream = -1;
	av_register_all();
	avformat_network_init();
	avcodec_register_all();
	avdevice_register_all();
}

CMediaSourceRtsp::~CMediaSourceRtsp()
{
	CloseInput();
	CloseOutput();
}

void CMediaSourceRtsp::CloseInput()
{
	if (inputContext) {
		avformat_close_input(&inputContext);
		inputContext = nullptr;
	}
}

void CMediaSourceRtsp::CloseOutput()
{
	if (outputContext) {
		av_write_trailer(outputContext);
		for (int i = 0; i < outputContext->nb_streams; i++)
		{
			if (outputContext->streams[i]->codec)
				avcodec_close(outputContext->streams[i]->codec);
		}
		avformat_close_input(&outputContext);
		outputContext = nullptr;
	}
}

void CMediaSourceRtsp::ActionOpenInput()
{
	while (1) {
		if (OpenInput(m_Input_rtsp, m_TcpOrudp))
			break;
		std::this_thread::sleep_for(std::chrono::milliseconds(2000));
	}
}

void CMediaSourceRtsp::ActionOpenOutput()
{
	while (1) {
		if (OpenOutput(m_Output_rtsp))
			break;
		std::this_thread::sleep_for(std::chrono::milliseconds(5000));
	}
}

void CMediaSourceRtsp::CheckOutputNet()
{
	m_NetCount++;
	if (m_NetCount > 500) {
		if (outputContext) {
			m_NetCount = 0;
			m_nWriteflag = false;
			CloseOutput();
			ActionOpenOutput();
		}
	}
}

void CMediaSourceRtsp::ResetConnect()
{
	if (inputContext) {
		CloseInput();
		CloseOutput();
	}
	//这里会一直阻塞，知道视频源和输出源正常连接
	ActionOpenInput();
	ActionOpenOutput();
}

bool CMediaSourceRtsp::OpenInput(std::string inputUrl, std::string transport)
{
	bool result = false;

	inputContext = avformat_alloc_context();
	AVDictionary* options = nullptr;

	if (transport == "tcp") {

		av_dict_set(&options, "rtsp_transport", "tcp", 0);
		av_dict_set(&options, "stimeout", "2000000", 0);

		if (0 > avformat_open_input(&inputContext, inputUrl.c_str(), nullptr, &options)) {

			OutputDebugStringA("failed open input file1");
			
			return result;
		}

	} else {

		av_dict_set(&options, "rtsp_transport", "udp", 0);
		av_dict_set(&options, "stimeout", "2000000", 0);
		if (0 > avformat_open_input(&inputContext, inputUrl.c_str(), nullptr, &options)) {

			OutputDebugStringA("failed open input file1");
			
			return result;
		}
	}

	int ret = avformat_find_stream_info(inputContext, nullptr);
	if (ret < 0)
	{
		OutputDebugStringA("Find input file stream inform failed");
		return result;
	}
	m_nVideoStream = av_find_best_stream(inputContext, AVMEDIA_TYPE_VIDEO, -1, -1, NULL, 0);
	if (-1 == m_nVideoStream) {
		OutputDebugStringA("failed av_find_best_stream");
		return result;
	}
	return true;
}
int  CMediaSourceRtsp::InterruptCb(void* ctx)
{
	CMediaSourceRtsp* media_source = (CMediaSourceRtsp*)ctx;

	if (media_source->m_nWriteflag){
		return 0;
	}
	int64_t localTime = av_gettime();
	
	//1s = 1000*1000微妙
	int times = (localTime % (media_source->m_nStartTime))/(1000*1000);
	if (times > media_source->m_timeout)
	{
		return 1;
	}
	return 0;
}

int CMediaSourceRtsp::WritePacket(std::shared_ptr<AVPacket> packet)
{
	auto inputStream = inputContext->streams[packet->stream_index];
	auto outputStream = outputContext->streams[packet->stream_index];
	av_packet_rescale_ts(packet.get(), inputStream->time_base, outputStream->time_base);
	return av_interleaved_write_frame(outputContext, packet.get());
}

std::shared_ptr<AVPacket> CMediaSourceRtsp::ReadPacket()
{
	std::shared_ptr<AVPacket> packet(static_cast<AVPacket*>(av_malloc(sizeof(AVPacket))), [&](AVPacket* p) { av_packet_free(&p); av_freep(&p); });
	av_init_packet(packet.get());
	int ret = av_read_frame(inputContext, packet.get());
	if (ret >= 0)
	{
		return packet;
	}
	else
	{
		return nullptr;
	}
}

bool CMediaSourceRtsp::OpenOutput(std::string outUrl)
{
	bool result = false;
	int ret = avformat_alloc_output_context2(&outputContext, nullptr, "rtsp", outUrl.c_str());
	if (ret < 0)
	{
		OutputDebugStringA("open output context failed");
		
		return result;
	}
	//设置超时回调
	outputContext->interrupt_callback.callback = InterruptCb;    // 设置超时回调
	outputContext->interrupt_callback.opaque = this;


	//for (int i = 0; i < inputContext->nb_streams; i++)
	//{
	//	AVStream* stream = avformat_new_stream(outputContext, inputContext->streams[i]->codec->codec);
	//	ret = avcodec_copy_context(stream->codec, inputContext->streams[i]->codec);
	//	if (ret < 0)
	//	{
	//		av_log(NULL, AV_LOG_ERROR, "copy coddec context failed");
	//		return result;
	//	}
	//}
	AVStream* stream = avformat_new_stream(outputContext, inputContext->streams[m_nVideoStream]->codec->codec);
	avcodec_copy_context(stream->codec, inputContext->streams[m_nVideoStream]->codec);
	if (ret < 0)
	{
		OutputDebugStringA("copy coddec context failed");
		return result;
	}

	//需要留意的是发输出是rtsp流时，是不需要调用avio_open，如果你调用了返回值是<0，所以这里m_pOutfmt_ctx->oformat->flags & AVFMT_NOFILE条件判断
	if (!(outputContext->oformat->flags & AVFMT_NOFILE)) {
		if (0 > avio_open(&outputContext->pb, outUrl.c_str(), AVIO_FLAG_WRITE)) {
			OutputDebugStringA("open avio failed");
			return result;
		}
	}

	m_nStartTime = av_gettime();
	ret = avformat_write_header(outputContext, nullptr);
	if (ret < 0)
	{
		OutputDebugStringA("format write header failed");
		
		return result;
	}
	m_nWriteflag = true;
	return true;
}