FLV: Add keyframes metadata

BililiveRecorder.Core/Recording/FlvProcessingContextWriterWithFileWriterFactory.cs

@@ -16,8 +16,11 @@ namespace BililiveRecorder.Core.Recording
         }
 
         public IFlvProcessingContextWriter CreateWriter(IFlvWriterTargetProvider targetProvider) =>
-            new FlvProcessingContextWriter(new FlvTagFileWriter(targetProvider,
+            new FlvProcessingContextWriter(
-                                                                this.serviceProvider.GetRequiredService<IMemoryStreamProvider>(),
+                tagWriter: new FlvTagFileWriter(targetProvider: targetProvider,
-                                                                this.serviceProvider.GetService<ILogger>()));
+                                                memoryStreamProvider: this.serviceProvider.GetRequiredService<IMemoryStreamProvider>(),
+                                                logger: this.serviceProvider.GetService<ILogger>()),
+                allowMissingHeader: false,
+                disableKeyframes: false);
     }
 }

BililiveRecorder.Flv/Pipeline/Rules/HandleNewScriptRule.cs

@@ -13,7 +13,6 @@ namespace BililiveRecorder.Flv.Pipeline.Rules
     /// </remarks>
     public class HandleNewScriptRule : ISimpleProcessingRule
     {
-        private static readonly ProcessingComment comment_other = new ProcessingComment(CommentType.Logging, "收到了非 onMetaData 的 Script Tag");
         private static readonly ProcessingComment comment_onmetadata = new ProcessingComment(CommentType.Logging, "收到了 onMetaData");
 
         public void Run(FlvProcessingContext context, Action next)
@@ -32,16 +31,13 @@ namespace BililiveRecorder.Flv.Pipeline.Rules
                     && data.Values[0] is ScriptDataString name
                     && name == "onMetaData")
                 {
-                    ScriptDataEcmaArray? value = data.Values[1] switch
+                    ScriptDataEcmaArray value = data.Values[1] switch
                     {
                         ScriptDataObject obj => obj,
                         ScriptDataEcmaArray arr => arr,
-                        _ => null
+                        _ => new ScriptDataEcmaArray()
                     };
 
-                    if (value is null)
-                        value = new ScriptDataEcmaArray();
-
                     context.AddComment(comment_onmetadata);
                     yield return PipelineNewFileAction.Instance;
                     yield return (new PipelineScriptAction(new Tag
@@ -56,7 +52,7 @@ namespace BililiveRecorder.Flv.Pipeline.Rules
                 }
                 else
                 {
-                    context.AddComment(comment_other);
+                    context.AddComment(new ProcessingComment(CommentType.Logging, "收到了非 onMetaData 的 Script Tag: " + (data?.ToJson() ?? "(null)")));
                 }
             }
             else

BililiveRecorder.Flv/StreamExtensions.cs

@@ -1,5 +1,6 @@
 using System;
 using System.IO;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 
@@ -101,6 +102,7 @@ namespace BililiveRecorder.Flv
             return b;
         }
 
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
         internal static void Write(this Stream stream, byte[] bytes) => stream.Write(bytes, 0, bytes.Length);
 
         internal static bool SequenceEqual(this Stream self, Stream? other)

BililiveRecorder.Flv/Writer/FlvProcessingContextWriter.cs

@@ -12,6 +12,7 @@ namespace BililiveRecorder.Flv.Writer
         private readonly SemaphoreSlim semaphoreSlim = new SemaphoreSlim(1, 1);
         private readonly IFlvTagWriter tagWriter;
         private readonly bool allowMissingHeader;
+        private readonly bool disableKeyframes;
         private bool disposedValue;
 
         private WriterState state = WriterState.EmptyFileOrNotOpen;
@@ -21,6 +22,7 @@ namespace BililiveRecorder.Flv.Writer
         private Tag? nextVideoHeaderTag = null;
 
         private ScriptTagBody? lastScriptBody = null;
+        private KeyframesScriptDataValue? keyframesScriptDataValue = null;
         private double lastDuration;
 
         public event EventHandler<FileClosedEventArgs>? FileClosed;
@@ -28,13 +30,11 @@ namespace BililiveRecorder.Flv.Writer
         public Action<ScriptTagBody>? BeforeScriptTagWrite { get; set; }
         public Action<ScriptTagBody>? BeforeScriptTagRewrite { get; set; }
 
-        public FlvProcessingContextWriter(IFlvTagWriter tagWriter) : this(tagWriter: tagWriter, allowMissingHeader: false)
+        public FlvProcessingContextWriter(IFlvTagWriter tagWriter, bool allowMissingHeader, bool disableKeyframes)
-        { }
-
-        public FlvProcessingContextWriter(IFlvTagWriter tagWriter, bool allowMissingHeader)
         {
             this.tagWriter = tagWriter ?? throw new ArgumentNullException(nameof(tagWriter));
             this.allowMissingHeader = allowMissingHeader;
+            this.disableKeyframes = disableKeyframes;
         }
 
         public async Task WriteAsync(FlvProcessingContext context)
@@ -147,7 +147,7 @@ namespace BililiveRecorder.Flv.Writer
             this.state = WriterState.BeforeScript;
         }
 
-        private async Task RewriteScriptTagImpl(double duration)
+        private async Task RewriteScriptTagImpl(double duration, bool updateKeyframes, double keyframeTime, double filePosition)
         {
             if (this.lastScriptBody is null)
                 return;
@@ -156,6 +156,9 @@ namespace BililiveRecorder.Flv.Writer
             if (value is not null)
                 value["duration"] = (ScriptDataNumber)duration;
 
+            if (updateKeyframes && this.keyframesScriptDataValue is not null)
+                this.keyframesScriptDataValue.AddData(keyframeTime, filePosition);
+
             this.BeforeScriptTagRewrite?.Invoke(this.lastScriptBody);
 
             await this.tagWriter.OverwriteMetadata(this.lastScriptBody).ConfigureAwait(false);
@@ -173,8 +176,17 @@ namespace BililiveRecorder.Flv.Writer
 
             var value = this.lastScriptBody.GetMetadataValue();
             if (value is not null)
+            {
                 value["duration"] = (ScriptDataNumber)0;
 
+                if (!this.disableKeyframes)
+                {
+                    var kfv = new KeyframesScriptDataValue();
+                    value["keyframes"] = kfv;
+                    this.keyframesScriptDataValue = kfv;
+                }
+            }
+
             this.BeforeScriptTagWrite?.Invoke(this.lastScriptBody);
 
             await this.tagWriter.WriteTag(this.nextScriptTag).ConfigureAwait(false);
@@ -228,12 +240,16 @@ namespace BililiveRecorder.Flv.Writer
                     throw new InvalidOperationException($"Can't write data tag with current state ({this.state})");
             }
 
-            foreach (var tag in dataAction.Tags)
+            var pos = this.tagWriter.FileSize;
+            var tags = dataAction.Tags;
+            var firstTag = tags[0];
+            var duration = tags[tags.Count - 1].Timestamp / 1000d;
+            this.lastDuration = duration;
+
+            foreach (var tag in tags)
                 await this.tagWriter.WriteTag(tag).ConfigureAwait(false);
 
-            var duration = dataAction.Tags[dataAction.Tags.Count - 1].Timestamp / 1000d;
+            await this.RewriteScriptTagImpl(duration, firstTag.IsKeyframeData(), firstTag.Timestamp, pos).ConfigureAwait(false);
-            this.lastDuration = duration;
-            await this.RewriteScriptTagImpl(duration).ConfigureAwait(false);
         }
 
         private async Task WriteEndTag(PipelineEndAction endAction)

BililiveRecorder.Flv/Writer/KeyframesScriptDataValue.cs

@@ -0,0 +1,172 @@
+using System;
+using System.Buffers.Binary;
+using System.Collections.Generic;
+using System.IO;
+using System.Runtime.CompilerServices;
+using System.Text;
+using BililiveRecorder.Flv.Amf;
+
+namespace BililiveRecorder.Flv.Writer
+{
+    /// <summary>
+    /// 输出时用于生成 OnMetaData 中的 keyframes。直接建 object tree 再序列化太太太太慢了
+    /// </summary>
+    internal class KeyframesScriptDataValue : IScriptDataValue
+    {
+        /* 
+         * 最少能保存大约 6300 * 2 second = 3.5 hour 的关键帧索引
+         * 如果以 5 秒计算则 6300 * 5 second = 8.75 hour
+         */
+
+        private const int MaxDataCount = 6300;
+        private const double MinInterval = 1900;
+
+        private const string Keyframes = "keyframes";
+        private const string Times = "times";
+        private const string FilePositions = "filepositions";
+        private const string Spacer = "spacer";
+
+        private static readonly byte[] EndBytes = new byte[] { 0, 0, 9 };
+
+        private static readonly byte[] TimesBytes = Encoding.UTF8.GetBytes(Times);
+        private static readonly byte[] FilePositionsBytes = Encoding.UTF8.GetBytes(FilePositions);
+        private static readonly byte[] SpacerBytes = Encoding.UTF8.GetBytes(Spacer);
+
+        public ScriptDataType Type => ScriptDataType.Object;
+
+        private readonly List<Data> KeyframesData = new();
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        public void AddData(double time, double filePosition)
+        {
+            var keyframesData = this.KeyframesData;
+            if (keyframesData.Count < MaxDataCount && (keyframesData.Count == 0 || ((time - keyframesData[keyframesData.Count - 1].Time) > MinInterval)))
+            {
+                keyframesData.Add(new Data(time: time, filePosition: filePosition));
+            }
+        }
+
+        /// <summary>
+        /// <see cref="ScriptDataObject.WriteTo(Stream)"/>
+        /// <see cref="ScriptDataStrictArray.WriteTo(Stream)"/>
+        /// <see cref="ScriptDataNumber.WriteTo(Stream)"/>
+        /// </summary>
+        /// <param name="stream"></param>
+        public void WriteTo(Stream stream)
+        {
+            stream.WriteByte((byte)this.Type);
+
+            var keyframesData = this.KeyframesData;
+            var buffer = new byte[sizeof(double)];
+
+            {
+                // key
+                WriteKey(stream, TimesBytes);
+
+                // array
+                WriteStrictArray(stream, (uint)keyframesData.Count);
+
+                // value
+                for (var i = 0; i < keyframesData.Count; i++)
+                {
+                    stream.WriteByte((byte)ScriptDataType.Number);
+                    BinaryPrimitives.WriteInt64BigEndian(buffer, BitConverter.DoubleToInt64Bits(keyframesData[i].Time));
+                    stream.Write(buffer);
+                }
+            }
+
+            {
+                // key
+                WriteKey(stream, FilePositionsBytes);
+
+                // array
+                WriteStrictArray(stream, (uint)keyframesData.Count);
+
+                // value
+                for (var i = 0; i < keyframesData.Count; i++)
+                {
+                    stream.WriteByte((byte)ScriptDataType.Number);
+                    BinaryPrimitives.WriteInt64BigEndian(buffer, BitConverter.DoubleToInt64Bits(keyframesData[i].FilePosition));
+                    stream.Write(buffer);
+                }
+            }
+
+            {
+                // key
+                WriteKey(stream, SpacerBytes);
+
+                // array
+                var count = 2u * (uint)(MaxDataCount - keyframesData.Count);
+                WriteStrictArray(stream, count);
+
+                // value
+                BinaryPrimitives.WriteInt64BigEndian(buffer, BitConverter.DoubleToInt64Bits(double.NaN));
+                for (var i = 0; i < count; i++)
+                {
+                    stream.WriteByte((byte)ScriptDataType.Number);
+                    stream.Write(buffer);
+                }
+            }
+
+            stream.Write(EndBytes);
+        }
+
+        public readonly struct Data
+        {
+            public readonly double Time;
+            public readonly double FilePosition;
+
+            public Data(double time, double filePosition)
+            {
+                this.Time = time;
+                this.FilePosition = filePosition;
+            }
+        }
+
+        /// <summary>
+        /// <see cref="ScriptDataStrictArray.WriteTo(Stream)"/>
+        /// </summary>
+        /// <param name="stream"></param>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static unsafe void WriteStrictArray(Stream stream, uint count)
+        {
+            stream.WriteByte((byte)ScriptDataType.StrictArray);
+
+            var buffer = new byte[sizeof(uint)];
+            BinaryPrimitives.WriteUInt32BigEndian(buffer, count);
+            stream.Write(buffer);
+        }
+
+        /// <summary>
+        /// <see cref="ScriptDataObject.WriteTo(Stream)"/>
+        /// </summary>
+        /// <param name="stream"></param>
+        /// <param name="key"></param>
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static void WriteKey(Stream stream, string key)
+        {
+            var bytes = Encoding.UTF8.GetBytes(key);
+            if (bytes.Length > ushort.MaxValue)
+                throw new AmfException($"Cannot write more than {ushort.MaxValue} into ScriptDataString");
+
+            var buffer = new byte[sizeof(ushort)];
+            BinaryPrimitives.WriteUInt16BigEndian(buffer, (ushort)bytes.Length);
+
+            stream.Write(buffer);
+            stream.Write(bytes);
+        }
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static void WriteKey(Stream stream, byte[] bytes)
+        {
+            //if (bytes.Length > ushort.MaxValue)
+            //    throw new AmfException($"Cannot write more than {ushort.MaxValue} into ScriptDataString");
+
+            var buffer = new byte[sizeof(ushort)];
+            BinaryPrimitives.WriteUInt16BigEndian(buffer, (ushort)bytes.Length);
+
+            stream.Write(buffer);
+            stream.Write(bytes);
+        }
+    }
+}

BililiveRecorder.ToolBox/Commands/Analyze.cs

@@ -106,7 +106,7 @@ namespace BililiveRecorder.ToolBox.Commands
 
                 // Pipeline
                 using var grouping = new TagGroupReader(tagReader);
-                using var writer = new FlvProcessingContextWriter(tagWriter: tagWriter, allowMissingHeader: true);
+                using var writer = new FlvProcessingContextWriter(tagWriter: tagWriter, allowMissingHeader: true, disableKeyframes: true);
                 var statsRule = new StatsRule();
                 var pipeline = new ProcessingPipelineBuilder(new ServiceCollection().BuildServiceProvider()).Add(statsRule).AddDefault().AddRemoveFillerData().Build();
 

BililiveRecorder.ToolBox/Commands/Fix.cs

@@ -127,7 +127,7 @@ namespace BililiveRecorder.ToolBox.Commands
 
                 // Pipeline
                 using var grouping = new TagGroupReader(tagReader);
-                using var writer = new FlvProcessingContextWriter(tagWriter: tagWriter, allowMissingHeader: true);
+                using var writer = new FlvProcessingContextWriter(tagWriter: tagWriter, allowMissingHeader: true, disableKeyframes: false);
                 var statsRule = new StatsRule();
                 var pipeline = new ProcessingPipelineBuilder(new ServiceCollection().BuildServiceProvider()).Add(statsRule).AddDefault().AddRemoveFillerData().Build();
 

test/BililiveRecorder.Flv.RuleTests/Integrated/TestBase.cs

@@ -40,7 +40,7 @@ namespace BililiveRecorder.Flv.RuleTests.Integrated
 
         protected async Task RunPipeline(ITagGroupReader reader, IFlvTagWriter output, List<ProcessingComment> comments)
         {
-            var writer = new FlvProcessingContextWriter(tagWriter: output, allowMissingHeader: true);
+            var writer = new FlvProcessingContextWriter(tagWriter: output, allowMissingHeader: true, disableKeyframes: true);
             var session = new Dictionary<object, object?>();
             var context = new FlvProcessingContext();
             var pipeline = this.BuildPipeline();

test/BililiveRecorder.Flv.UnitTests/Grouping/GroupingTest.cs

@@ -81,7 +81,7 @@ namespace BililiveRecorder.Flv.UnitTests.Grouping
             var sp = new ServiceCollection().BuildServiceProvider();
             var pipeline = new ProcessingPipelineBuilder(sp).AddDefault().AddRemoveFillerData().Build();
 
-            using var writer = new FlvProcessingContextWriter(new FlvTagFileWriter(new TestOutputProvider(), new TestRecyclableMemoryStreamProvider(), null));
+            using var writer = new FlvProcessingContextWriter(tagWriter: new FlvTagFileWriter(new TestOutputProvider(), new TestRecyclableMemoryStreamProvider(), null), allowMissingHeader: false, disableKeyframes: true);
 
             while (true)
             {