RTMP/Assets/RTMPPublisher.cs

using System;
using System.Collections;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Threading;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using UnityEngine;
using UnityEngine.Experimental.Rendering;
using UnityEngine.Profiling;
using UnityEngine.Rendering;
using Debug = UnityEngine.Debug;
using Object = UnityEngine.Object;

namespace ZC
{
    public class RTMPPublisher : MonoBehaviour
    {
        Thread _thread;
        private Timer _timer;
        Process _process;
        private RectInt _rectInt;
        private RTMPConfig _config;
        private const int _rectIntWidth = 1920;
        private const int _rectIntHeight = 1080;


        [SerializeField] private Camera _camera;

        private Texture2D _output;

        private volatile int _sendData;

        private AsyncGPUReadbackRequest _asyncGPUReadbackRequest;
        private bool _isRequesting;


        private byte[] _bitmapArray;
        private ConcurrentQueue<NativeArray<byte>> _textureDataPool;
        private ConcurrentDictionary<uint, NativeArray<byte>> _doneTextureDatas;
        private int _targetTextureWidth;
        private int _targetTextureHeight;
        private uint _frameIndex;
        private uint _doneFrameIndex;
        private int _frameCount;

        public void SetCamera(Camera camera)
        {
            this._camera = camera;
        }

        // Start is called before the first frame update
        void Start()
        {
            _config = Object.FindAnyObjectByType<RTMPConfig>();
            if (_config == null)
            {
                Debug.LogError("This need the config, plz create a GameObject with RTMPConfig");
                return;
            }

            _textureDataPool = new ConcurrentQueue<NativeArray<byte>>();
            _doneTextureDatas = new ConcurrentDictionary<uint, NativeArray<byte>>();
            _config.Load();
            _timer = new Timer(Tick, null, TimeSpan.FromMilliseconds(0), TimeSpan.FromMilliseconds(16));
            Setup(Camera.main);
        }

        /// <summary>
        /// 其他线程tick
        /// </summary>
        /// <param name="state"></param>
        private void Tick(object state)
        {
            try
            {
                if (_doneTextureDatas.TryRemove(this._doneFrameIndex, out var result))
                {
                    var textureData = result;
                    Profiler.BeginSample("Conversion");
                    Bitmap.EncodeToBitmap(textureData, _bitmapArray, 0, textureData.Length, _targetTextureWidth, this._targetTextureHeight);
                    Profiler.EndSample();
                    Profiler.BeginSample("WriteBuffer");
                    _process.StandardInput.BaseStream.Write(_bitmapArray, 0, _bitmapArray.Length);
                    Profiler.EndSample();
                    _textureDataPool.Enqueue(result);
                    _ = _doneFrameIndex + 1 >= uint.MaxValue ? _doneFrameIndex = 0 : _doneFrameIndex++;
                }
            }
            catch (Exception e)
            {
                UnityEngine.Debug.LogError(e.ToString());
            }
        }

        private void CreateCaptureThread()
        {
            this._thread = new Thread(this.CreateThread);
            this._thread.Start();
        }


        private void OnDestroy()
        {
            Dispose();
        }

        private void Update()
        {
            this._frameCount = Time.frameCount;
            if (_sendData != 0)
            {
                if (_camera && !_isRequesting)
                {
                    var cameraActiveTexture = _camera.targetTexture;
                    RenderTexture.active = cameraActiveTexture;
                    if (!_textureDataPool.TryDequeue(out var nativeArray))
                    {
                        var targetTextureByteCount = this._camera.targetTexture.width * _camera.targetTexture.height * 3;
                        nativeArray = new NativeArray<byte>(targetTextureByteCount, Allocator.Persistent);
                    }

                    var frameIndex = this._frameIndex + 1 >= uint.MaxValue ? this._frameIndex = 0 : this._frameIndex++;
                    AsyncGPUReadback.RequestIntoNativeArray(ref nativeArray, cameraActiveTexture, 0, GraphicsFormat.B8G8R8_SRGB, request =>
                    {
                        if (request is { done: true, hasError: false })
                        {
                            this._doneTextureDatas[frameIndex] = nativeArray;
                        }
                        else
                        {
                            this._textureDataPool.Enqueue(nativeArray);
                            Debug.LogError($"Done:{request.done} Error:{request.hasError}");
                        }
                    });
                }
            }
        }


        private void DisposeCaptureThread()
        {
            foreach (var kv in this._doneTextureDatas)
            {
                kv.Value.Dispose();
            }

            _doneTextureDatas.Clear();

            foreach (var nativeArray in this._textureDataPool)
            {
                nativeArray.Dispose();
            }

            _textureDataPool.Clear();

            if (_process != null && !this._process.HasExited)
            {
                GenerateConsoleCtrlEvent(ConsoleCtrlEvent.CTRL_C, 0);

//                _process.StandardInput.Close();
                this._process.Kill();
                this._process.Close();
            }

            this._thread.Abort();
        }

        void CreateThread(object state)
        {
            _process = new Process();
            ProcessStartInfo processStartInfo = new ProcessStartInfo();
            processStartInfo.UseShellExecute = false;
            processStartInfo.CreateNoWindow = true;
            processStartInfo.RedirectStandardInput = true;
            processStartInfo.RedirectStandardOutput = true;
            processStartInfo.RedirectStandardError = true;
            processStartInfo.StandardOutputEncoding = System.Text.Encoding.UTF8;
            processStartInfo.StandardErrorEncoding = System.Text.Encoding.UTF8;
            processStartInfo.FileName = $"\"{_config.ffmpegPath}\"";
            processStartInfo.Arguments =
                $" -probesize 64 -thread_queue_size 5096 -fflags discardcorrupt -flags low_delay -analyzeduration 0 " +
                $" -rtbufsize 100M -f dshow -i audio=\"virtual-audio-capturer\" " +
                $" -f image2pipe -use_wallclock_as_timestamps 1 -r 60 -i - " +
                $" -loglevel info " +
                $" -map 0:a:0  -map 1:v:0 " +
                $" -c:a aac " +
                $" -c:v:0 libx264 -g 1 -bf 0 -speed 1x -max_delay 0 -vf scale={this._config.resolution} -preset:v ultrafast -tune:v zerolatency -crf 10 -pix_fmt yuv420p -strict -2 " +
                $" -f flv {this._config.server}{this._config.appName} -bf 0 ";
            Debug.Log(processStartInfo.Arguments);

            _process.StartInfo = processStartInfo;
            _process.EnableRaisingEvents = true;
            _process.OutputDataReceived += (s, e) => { Debug.Log(e.Data); };
            _process.ErrorDataReceived += (s, e) =>
            {
                if (!string.IsNullOrEmpty(e.Data) && e.Data.ToLower().Contains("error"))
                {
                    Debug.LogError(e.Data);
                }
                else
                {
                    Debug.Log(e.Data);
                }
            };
            _process.Start();
            Interlocked.Exchange(ref _sendData, 1);
            _process.BeginOutputReadLine();
            _process.BeginErrorReadLine();
            _process.WaitForExit();
            Debug.Log("Process exited!");
        }

        public void Setup(Camera camera)
        {
            if (!camera)
                throw new NullReferenceException("camera is null");
            if (!camera.targetTexture)
                throw new Exception($"The camera[{camera}]'s targetTexture is null;");
            _isRequesting = true;
            SetCamera(camera);
            _sendData = 0;
            Object.Destroy(_output);
            this._targetTextureWidth = camera.targetTexture.width;
            this._targetTextureHeight = camera.targetTexture.height;
            var targetTextureByteCount = this._targetTextureWidth * this._targetTextureHeight * 3;
            if (_bitmapArray == null || _bitmapArray.Length <
                Bitmap.FileHeaderSize + Bitmap.ImageHeaderSize + targetTextureByteCount)
            {
                _bitmapArray = new byte[Bitmap.FileHeaderSize + Bitmap.ImageHeaderSize + targetTextureByteCount];
            }

            _output = new Texture2D(this._targetTextureWidth, this._targetTextureHeight);

            _frameIndex = _doneFrameIndex = 0;
            CreateCaptureThread();

            StartCoroutine(this.Test());
        }


        IEnumerator Test()
        {
            yield return new WaitForSeconds(0.3f);
            _isRequesting = false;
        }

        public void Dispose()
        {
            try
            {
                this._isRequesting = true;
                _sendData = 0;
                _camera = null;
                Object.Destroy(_output);
                this.DisposeCaptureThread();
            }
            finally
            {
#if UNITY_EDITOR
#else
                _textureData.Dispose();
#endif
            }
        }

        [DllImport("kernel32.dll", SetLastError = true)]
        static extern bool GenerateConsoleCtrlEvent(ConsoleCtrlEvent sigevent, int dwProcessGroupId);

        public enum ConsoleCtrlEvent
        {
            CTRL_C = 0,
            CTRL_BREAK = 1,
            CTRL_CLOSE = 2,
            CTRL_LOGOFF = 5,
            CTRL_SHUTDOWN = 6
        }
    }
}