Stop encoding when the task is cancelled

This commit is contained in:
Sami Samhuri 2024-07-10 23:30:28 -07:00
parent f49cc722d4
commit fa1f39bb2b
No known key found for this signature in database
2 changed files with 87 additions and 16 deletions

View file

@ -44,6 +44,7 @@ actor SampleWriter {
private var audioInput: AVAssetWriterInput?
private var videoOutput: AVAssetReaderVideoCompositionOutput?
private var videoInput: AVAssetWriterInput?
private var isCancelled = false
init(
asset: sending AVAsset,
@ -90,15 +91,19 @@ actor SampleWriter {
}
func writeSamples() async throws {
try Task.checkCancellation()
progressContinuation?.yield(0.0)
writer.startWriting()
reader.startReading()
writer.startSession(atSourceTime: timeRange.start)
reader.startReading()
try Task.checkCancellation()
await encodeVideoTracks()
await encodeAudioTracks()
try Task.checkCancellation()
await encodeAudioTracks()
try Task.checkCancellation()
guard reader.status != .cancelled && writer.status != .cancelled else {
@ -173,34 +178,76 @@ actor SampleWriter {
self.videoInput = videoInput
}
func cancel() async {
isCancelled = true
}
// MARK: - Encoding
private func encodeAudioTracks() async {
// Don't do anything when we have no audio to encode.
guard audioInput != nil, audioOutput != nil else { return }
return await withCheckedContinuation { continuation in
self.audioInput!.requestMediaDataWhenReady(on: queue) {
let hasMoreSamples = self.assumeIsolated { _self in
_self.writeReadySamples(output: _self.audioOutput!, input: _self.audioInput!)
}
if !hasMoreSamples {
continuation.resume()
return await withTaskCancellationHandler {
await withCheckedContinuation { continuation in
self.audioInput!.requestMediaDataWhenReady(on: queue) {
self.assumeIsolated { _self in
guard !_self.isCancelled else {
log.debug("Cancelled while encoding audio")
_self.reader.cancelReading()
_self.writer.cancelWriting()
continuation.resume()
return
}
let hasMoreSamples = _self.writeReadySamples(
output: _self.audioOutput!,
input: _self.audioInput!
)
if !hasMoreSamples {
log.debug("Finished encoding audio")
continuation.resume()
}
}
}
}
} onCancel: {
log.debug("Task cancelled while encoding audio")
Task {
await self.cancel()
}
}
}
private func encodeVideoTracks() async {
return await withCheckedContinuation { continuation in
self.videoInput!.requestMediaDataWhenReady(on: queue) {
let hasMoreSamples = self.assumeIsolated { _self in
_self.writeReadySamples(output: _self.videoOutput!, input: _self.videoInput!)
}
if !hasMoreSamples {
continuation.resume()
return await withTaskCancellationHandler {
await withCheckedContinuation { continuation in
self.videoInput!.requestMediaDataWhenReady(on: queue) {
self.assumeIsolated { _self in
guard !_self.isCancelled else {
log.debug("Cancelled while encoding video")
_self.reader.cancelReading()
_self.writer.cancelWriting()
continuation.resume()
return
}
let hasMoreSamples = _self.writeReadySamples(
output: _self.videoOutput!,
input: _self.videoInput!
)
if !hasMoreSamples {
log.debug("Finished encoding video")
continuation.resume()
}
}
}
}
} onCancel: {
log.debug("Task cancelled while encoding video")
Task {
await self.cancel()
}
}
}

View file

@ -297,4 +297,28 @@ final class ExportSessionTests {
)
}
}
@Test func test_export_cancellation() async throws {
let sourceURL = resourceURL(named: "test-720p-h264-24fps", withExtension: "mov")
let destinationURL💥 = makeTemporaryURL()
let task = Task {
let sourceAsset = AVURLAsset(url: sourceURL, options: [
AVURLAssetPreferPreciseDurationAndTimingKey: true,
])
let subject = ExportSession()
try await subject.export(
asset: sourceAsset,
video: .codec(.h264, width: 1280, height: 720),
to: destinationURL💥.url,
as: .mov
)
Issue.record("Task should be cancelled long before we get here")
}
NSLog("Sleeping for 0.3s")
try await Task.sleep(for: .milliseconds(300))
NSLog("Cancelling task")
task.cancel()
try? await task.value // Wait for task to complete
NSLog("Task has finished executing")
}
}