Skip to content

Commit

Permalink
Clean up .iamf file when sequencing fails.
Browse files Browse the repository at this point in the history 
  - Context: `PickAndPlace` returns an error code when sequencing fails. In some circumstances it would leave a `.iamf` file even in error conditions.
  - Create the file immediately before writing to it (instead of in the constructor).
  - Use a separate internal function for `PickAndPlace` which writes to the file.
  - Remove the file when that function fails.
  - b/302470464: Backfill some coverage that `ObuSequencerIamf::PickAndPlace` actually writes files.

PiperOrigin-RevId: 690741424
  • Loading branch information
@jwcullen
jwcullen committed Nov 1, 2024
1 parent 05e8a14 commit 289fde2
Show file tree
Hide file tree
Showing 3 changed files with 162 additions and 31 deletions.
79 changes: 58 additions & 21 deletions iamf/cli/obu_sequencer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,14 @@

#include <algorithm>
#include <cstdint>
#include <filesystem>
#include <fstream>
#include <functional>
#include <ios>
#include <list>
#include <optional>
#include <string>
#include <system_error>
#include <utility>
#include <vector>

Expand Down Expand Up @@ -44,6 +49,9 @@ namespace iamf_tools {

namespace {

// Start with 64 KB.
constexpr int64_t kBufferStartSize = 65536;

template <typename KeyValueMap, typename KeyComparator>
std::vector<uint32_t> SortedKeys(const KeyValueMap& map,
const KeyComparator& comparator) {
Expand All @@ -56,6 +64,48 @@ std::vector<uint32_t> SortedKeys(const KeyValueMap& map,
return keys;
}

// The caller of this function should pre-seed the write buffer with Descriptor
// OBUs.
absl::Status WriteIaSequenceToFile(const std::string& iamf_filename,
bool include_temporal_delimiters,
const TemporalUnitMap& temporal_unit_map,
WriteBitBuffer& wb) {
std::optional<std::fstream> output_iamf;
if (!iamf_filename.empty()) {
output_iamf.emplace(iamf_filename, std::fstream::out | std::ios::binary);
}

// Write all Audio Frame and Parameter Block OBUs ordered by temporal unit.
int num_samples = 0;
for (const auto& temporal_unit : temporal_unit_map) {
// The temporal units will typically be the largest part of an IAMF
// sequence. Occasionally flush to buffer to avoid keeping it all in memory.
RETURN_IF_NOT_OK(wb.MaybeFlushIfCloseToCapacity(output_iamf));

RETURN_IF_NOT_OK(ObuSequencerBase::WriteTemporalUnit(
include_temporal_delimiters, temporal_unit.second, wb, num_samples));
}
LOG(INFO) << "Wrote " << temporal_unit_map.size()
<< " temporal units with a total of " << num_samples
<< " samples excluding padding.";

// Flush any unwritten bytes before exiting.
RETURN_IF_NOT_OK(wb.FlushAndWriteToFile(output_iamf));
return absl::OkStatus();
}

void MaybeRemoveFile(const std::string& filename) {
if (filename.empty()) {
return;
}
std::error_code error_code;
std::filesystem::remove(filename, error_code);
if (!error_code) {
// File clean up failed somehow. Just log the error and move on.
LOG(ERROR).WithPerror() << "Failed to remove " << filename;
}
}

} // namespace

absl::Status ObuSequencerBase::GenerateTemporalUnitMap(
Expand Down Expand Up @@ -290,10 +340,9 @@ absl::Status ObuSequencerIamf::PickAndPlace(
const std::list<AudioFrameWithData>& audio_frames,
const std::list<ParameterBlockWithData>& parameter_blocks,
const std::list<ArbitraryObu>& arbitrary_obus) {
// Write buffer. Let's start with 64 KB. The buffer will resize for larger
// OBUs if needed.
static const int64_t kBufferSize = 65536;
WriteBitBuffer wb(kBufferSize, leb_generator_);
// Seed with a reasonable starting size. It is arbitrary because
// `WriteBitBuffer`s automatically resize as needed.
WriteBitBuffer wb(kBufferStartSize, leb_generator_);

RETURN_IF_NOT_OK(ArbitraryObu::WriteObusWithHook(
ArbitraryObu::kInsertionHookBeforeDescriptors, arbitrary_obus, wb));
Expand All @@ -311,24 +360,12 @@ absl::Status ObuSequencerIamf::PickAndPlace(
RETURN_IF_NOT_OK(ObuSequencerBase::GenerateTemporalUnitMap(
audio_frames, parameter_blocks, arbitrary_obus, temporal_unit_map));

// Write all Audio Frame and Parameter Block OBUs ordered by temporal unit.
int num_samples = 0;
for (const auto& temporal_unit : temporal_unit_map) {
// The temporal units will typically be the largest part of an IAMF
// sequence. Occasionally flush to buffer to avoid keeping it all in memory.
RETURN_IF_NOT_OK(wb.MaybeFlushIfCloseToCapacity(output_iamf_));

RETURN_IF_NOT_OK(ObuSequencerBase::WriteTemporalUnit(
include_temporal_delimiters_, temporal_unit.second, wb, num_samples));
const auto write_status = WriteIaSequenceToFile(
iamf_filename_, include_temporal_delimiters_, temporal_unit_map, wb);
if (!write_status.ok()) {
MaybeRemoveFile(iamf_filename_);
}
LOG(INFO) << "Wrote " << temporal_unit_map.size()
<< " temporal units with a total of " << num_samples
<< " samples excluding padding.";

// Flush any unwritten bytes before exiting.
RETURN_IF_NOT_OK(wb.FlushAndWriteToFile(output_iamf_));

return absl::OkStatus();
return write_status;
}

} // namespace iamf_tools
10 changes: 2 additions & 8 deletions iamf/cli/obu_sequencer.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,9 +13,7 @@
#define CLI_OBU_SEQUENCER_H_

#include <cstdint>
#include <fstream>
#include <list>
#include <optional>
#include <string>
#include <vector>

Expand Down Expand Up @@ -150,11 +148,7 @@ class ObuSequencerIamf : public ObuSequencerBase {
bool include_temporal_delimiters,
const LebGenerator& leb_generator)
: ObuSequencerBase(leb_generator),
output_iamf_(
iamf_filename.empty()
? std::nullopt
: std::make_optional<std::fstream>(
iamf_filename, std::fstream::out | std::fstream::binary)),
iamf_filename_(iamf_filename),
include_temporal_delimiters_(include_temporal_delimiters) {}

~ObuSequencerIamf() override = default;
Expand All @@ -180,7 +174,7 @@ class ObuSequencerIamf : public ObuSequencerBase {
const std::list<ArbitraryObu>& arbitrary_obus) override;

private:
std::optional<std::fstream> output_iamf_;
const std::string iamf_filename_;
const bool include_temporal_delimiters_;
};

Expand Down
104 changes: 102 additions & 2 deletions iamf/cli/tests/obu_sequencer_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -51,15 +51,15 @@ namespace {
using ::absl_testing::IsOk;

constexpr DecodedUleb128 kCodecConfigId = 1;
const uint32_t kSampleRate = 48000;
constexpr uint32_t kSampleRate = 48000;
constexpr DecodedUleb128 kFirstAudioElementId = 1;
constexpr DecodedUleb128 kSecondAudioElementId = 2;
constexpr DecodedUleb128 kFirstSubstreamId = 1;
constexpr DecodedUleb128 kSecondSubstreamId = 2;
constexpr DecodedUleb128 kFirstMixPresentationId = 100;
constexpr DecodedUleb128 kFirstDemixingParameterId = 998;
constexpr DecodedUleb128 kCommonMixGainParameterId = 999;
const uint32_t kCommonMixGainParameterRate = kSampleRate;
constexpr uint32_t kCommonMixGainParameterRate = kSampleRate;

constexpr absl::string_view kOmitOutputIamfFile = "";
constexpr bool kIncludeTemporalDelimiters = true;
Expand Down Expand Up @@ -446,6 +446,21 @@ class ObuSequencerTest : public ::testing::Test {
ASSERT_FALSE(mix_presentation_obus_.empty());
}

void InitObusForOneFrameIaSequence() {
ia_sequence_header_obu_.emplace(ObuHeader(), IASequenceHeaderObu::kIaCode,
ProfileVersion::kIamfSimpleProfile,
ProfileVersion::kIamfSimpleProfile);
per_id_metadata_ =
CreatePerIdMetadataForDemixing(kFirstDemixingParameterId);
InitializeOneParameterBlockAndOneAudioFrame(
per_id_metadata_, parameter_blocks_, audio_frames_, codec_config_obus_,
audio_elements_);
AddMixPresentationObuWithAudioElementIds(
kFirstMixPresentationId, {audio_elements_.begin()->first},
kCommonMixGainParameterId, kCommonMixGainParameterRate,
mix_presentation_obus_);
}

void ValidateWriteDescriptorObuSequence(
const std::list<const ObuBase*>& expected_sequence) {
WriteBitBuffer expected_wb(128);
Expand All @@ -469,6 +484,11 @@ class ObuSequencerTest : public ::testing::Test {
absl::flat_hash_map<uint32_t, CodecConfigObu> codec_config_obus_;
absl::flat_hash_map<uint32_t, AudioElementWithData> audio_elements_;
std::list<MixPresentationObu> mix_presentation_obus_;

PerIdParameterMetadata per_id_metadata_;
std::list<ParameterBlockWithData> parameter_blocks_;
std::list<AudioFrameWithData> audio_frames_;

std::list<ArbitraryObu> arbitrary_obus_;
};

Expand Down Expand Up @@ -739,5 +759,85 @@ TEST(ObuSequencerIamf, PickAndPlaceSucceedsWithEmptyOutputFile) {
IsOk());
}

TEST_F(ObuSequencerTest, PickAndPlaceCreatesFileWithOneFrameIaSequence) {
const std::string kOutputIamfFilename = GetAndCleanupOutputFileName(".iamf");
InitObusForOneFrameIaSequence();
ObuSequencerIamf sequencer(kOutputIamfFilename,
kDoNotIncludeTemporalDelimiters,
*LebGenerator::Create());

ASSERT_THAT(
sequencer.PickAndPlace(*ia_sequence_header_obu_, codec_config_obus_,
audio_elements_, mix_presentation_obus_,
audio_frames_, parameter_blocks_, arbitrary_obus_),
IsOk());

EXPECT_TRUE(std::filesystem::exists(kOutputIamfFilename));
}

TEST_F(ObuSequencerTest, PickAndPlaceLeavesNoFileWhenDescriptorsAreInvalid) {
constexpr uint32_t kInvalidIaCode = IASequenceHeaderObu::kIaCode + 1;
const std::string kOutputIamfFilename = GetAndCleanupOutputFileName(".iamf");
InitObusForOneFrameIaSequence();
// Overwrite the IA Sequence Header with an invalid one.
ia_sequence_header_obu_ = IASequenceHeaderObu(
ObuHeader(), kInvalidIaCode, ProfileVersion::kIamfSimpleProfile,
ProfileVersion::kIamfSimpleProfile);
ObuSequencerIamf sequencer(kOutputIamfFilename,
kDoNotIncludeTemporalDelimiters,
*LebGenerator::Create());

ASSERT_FALSE(sequencer
.PickAndPlace(*ia_sequence_header_obu_, codec_config_obus_,
audio_elements_, mix_presentation_obus_,
audio_frames_, parameter_blocks_,
arbitrary_obus_)
.ok());

EXPECT_FALSE(std::filesystem::exists(kOutputIamfFilename));
}

TEST_F(ObuSequencerTest, PickAndPlaceLeavesNoFileWhenTemporalUnitsAreInvalid) {
constexpr bool kInvalidateTemporalUnit = true;
const std::string kOutputIamfFilename = GetAndCleanupOutputFileName(".iamf");
InitObusForOneFrameIaSequence();
arbitrary_obus_.emplace_back(
ArbitraryObu(kObuIaReserved25, ObuHeader(), {},
ArbitraryObu::kInsertionHookAfterAudioFramesAtTick, 0,
kInvalidateTemporalUnit));
ObuSequencerIamf sequencer(kOutputIamfFilename,
kDoNotIncludeTemporalDelimiters,
*LebGenerator::Create());

ASSERT_FALSE(sequencer
.PickAndPlace(*ia_sequence_header_obu_, codec_config_obus_,
audio_elements_, mix_presentation_obus_,
audio_frames_, parameter_blocks_,
arbitrary_obus_)
.ok());

EXPECT_FALSE(std::filesystem::exists(kOutputIamfFilename));
}

TEST_F(ObuSequencerTest,
PickAndPlaceOnInvalidTemporalUnitFailsWhenOutputFileIsOmitted) {
constexpr bool kInvalidateTemporalUnit = true;
InitObusForOneFrameIaSequence();
arbitrary_obus_.emplace_back(
ArbitraryObu(kObuIaReserved25, ObuHeader(), {},
ArbitraryObu::kInsertionHookAfterAudioFramesAtTick, 0,
kInvalidateTemporalUnit));
ObuSequencerIamf sequencer(std::string(kOmitOutputIamfFile),
kDoNotIncludeTemporalDelimiters,
*LebGenerator::Create());

ASSERT_FALSE(sequencer
.PickAndPlace(*ia_sequence_header_obu_, codec_config_obus_,
audio_elements_, mix_presentation_obus_,
audio_frames_, parameter_blocks_,
arbitrary_obus_)
.ok());
}

} // namespace
} // namespace iamf_tools

0 comments on commit 289fde2

Please sign in to comment.