From a13531321e7177a34d77a3c2b36afc93c263be73 Mon Sep 17 00:00:00 2001 From: Alex Gaynor Date: Fri, 18 May 2018 18:59:00 -0400 Subject: [PATCH] Bug 1456189 - Simplify BufferList::Extract to make the lifetimes clearer. r=froydnj --- mfbt/BufferList.h | 133 ++++++++++++++++++++++------------ mfbt/tests/TestBufferList.cpp | 33 ++++++++- 2 files changed, 117 insertions(+), 49 deletions(-) diff --git a/mfbt/BufferList.h b/mfbt/BufferList.h index 62ab540df0fb..a2e7aac32a9f 100644 --- a/mfbt/BufferList.h +++ b/mfbt/BufferList.h @@ -9,6 +9,7 @@ #include #include "mozilla/AllocPolicy.h" +#include "mozilla/Maybe.h" #include "mozilla/MemoryReporting.h" #include "mozilla/Move.h" #include "mozilla/ScopeExit.h" @@ -538,61 +539,101 @@ BufferList::Extract(IterImpl& aIter, size_t aSize, bool* aSuccess) MOZ_ASSERT(aSize % kSegmentAlignment == 0); MOZ_ASSERT(intptr_t(aIter.mData) % kSegmentAlignment == 0); - IterImpl iter = aIter; - size_t size = aSize; - size_t toCopy = std::min(size, aIter.RemainingInSegment()); - MOZ_ASSERT(toCopy % kSegmentAlignment == 0); - - BufferList result(0, toCopy, mStandardCapacity); - BufferList error(0, 0, mStandardCapacity); - - // Copy the head - if (!result.WriteBytes(aIter.mData, toCopy)) { + auto failure = [this, aSuccess]() { *aSuccess = false; - return error; - } - iter.Advance(*this, toCopy); - size -= toCopy; + return BufferList(0, 0, mStandardCapacity); + }; - // Move segments to result - auto resultGuard = MakeScopeExit([&] { - *aSuccess = false; - result.mSegments.erase(result.mSegments.begin()+1, result.mSegments.end()); - }); - - size_t movedSize = 0; - uintptr_t toRemoveStart = iter.mSegment; - uintptr_t toRemoveEnd = iter.mSegment; - while (!iter.Done() && - !iter.HasRoomFor(size)) { - if (!result.mSegments.append(Segment(mSegments[iter.mSegment].mData, - mSegments[iter.mSegment].mSize, - mSegments[iter.mSegment].mCapacity))) { - return error; + // Number of segments we'll need to copy data from to satisfy the request. + size_t segmentsNeeded = 0; + // If this is None then the last segment is a full segment, otherwise we need + // to copy this many bytes. + Maybe lastSegmentSize; + { + // Copy of the iterator to walk the BufferList and see how many segments we + // need to copy. + IterImpl iter = aIter; + size_t remaining = aSize; + while (!iter.Done() && remaining && + remaining >= iter.RemainingInSegment()) { + remaining -= iter.RemainingInSegment(); + iter.Advance(*this, iter.RemainingInSegment()); + segmentsNeeded++; } - movedSize += iter.RemainingInSegment(); - size -= iter.RemainingInSegment(); - toRemoveEnd++; - iter.Advance(*this, iter.RemainingInSegment()); - } - if (size) { - if (!iter.HasRoomFor(size) || - !result.WriteBytes(iter.Data(), size)) { - return error; + if (remaining) { + if (iter.Done()) { + // We reached the end of the BufferList and there wasn't enough data to + // satisfy the request. + return failure(); + } + lastSegmentSize.emplace(remaining); + // The last block also counts as a segment. This makes the conditionals + // on segmentsNeeded work in the rest of the function. + segmentsNeeded++; } - iter.Advance(*this, size); } - mSegments.erase(mSegments.begin() + toRemoveStart, mSegments.begin() + toRemoveEnd); - mSize -= movedSize; - aIter.mSegment = iter.mSegment - (toRemoveEnd - toRemoveStart); - aIter.mData = iter.mData; - aIter.mDataEnd = iter.mDataEnd; - MOZ_ASSERT(aIter.mDataEnd == mSegments[aIter.mSegment].End()); + BufferList result(0, 0, mStandardCapacity); + if (!result.mSegments.reserve(segmentsNeeded + lastSegmentSize.isSome())) { + return failure(); + } + + // Copy the first segment, it's special because we can't just steal the + // entire Segment struct from this->mSegments. + size_t firstSegmentSize = std::min(aSize, aIter.RemainingInSegment()); + if (!result.WriteBytes(aIter.Data(), firstSegmentSize)) { + return failure(); + } + aIter.Advance(*this, firstSegmentSize); + segmentsNeeded--; + + // The entirety of the request wasn't in the first segment, now copy the + // rest. + if (segmentsNeeded) { + char* finalSegment = nullptr; + // Pre-allocate the final segment so that if this fails, we return before + // we delete the elements from |this->mSegments|. + if (lastSegmentSize.isSome()) { + MOZ_RELEASE_ASSERT(mStandardCapacity >= *lastSegmentSize); + finalSegment = this->template pod_malloc(mStandardCapacity); + if (!finalSegment) { + return failure(); + } + } + + size_t copyStart = aIter.mSegment; + // Copy segments from this over to the result and remove them from our + // storage. Not needed if the only segment we need to copy is the last + // partial one. + size_t segmentsToCopy = segmentsNeeded - lastSegmentSize.isSome(); + for (size_t i = 0; i < segmentsToCopy; ++i) { + result.mSegments.infallibleAppend( + Segment(mSegments[aIter.mSegment].mData, + mSegments[aIter.mSegment].mSize, + mSegments[aIter.mSegment].mCapacity)); + aIter.Advance(*this, aIter.RemainingInSegment()); + } + MOZ_RELEASE_ASSERT(aIter.mSegment == copyStart + segmentsToCopy); + mSegments.erase(mSegments.begin() + copyStart, + mSegments.begin() + copyStart + segmentsToCopy); + + // Reset the iter's position for what we just deleted. + aIter.mSegment -= segmentsToCopy; + + if (lastSegmentSize.isSome()) { + // We called reserve() on result.mSegments so infallibleAppend is safe. + result.mSegments.infallibleAppend( + Segment(finalSegment, 0, mStandardCapacity)); + bool r = result.WriteBytes(aIter.Data(), *lastSegmentSize); + MOZ_RELEASE_ASSERT(r); + aIter.Advance(*this, *lastSegmentSize); + } + } + + mSize -= aSize; result.mSize = aSize; - resultGuard.release(); *aSuccess = true; return result; } diff --git a/mfbt/tests/TestBufferList.cpp b/mfbt/tests/TestBufferList.cpp index cccaac021b4a..207fa106f556 100644 --- a/mfbt/tests/TestBufferList.cpp +++ b/mfbt/tests/TestBufferList.cpp @@ -245,12 +245,39 @@ int main(void) BufferList bl3 = bl.Extract(iter, kExtractOverSize, &success); MOZ_RELEASE_ASSERT(!success); - MOZ_RELEASE_ASSERT(iter.AdvanceAcrossSegments(bl, kSmallWrite * 3 - kExtractSize - kExtractStart)); - MOZ_RELEASE_ASSERT(iter.Done()); - iter = bl2.Iter(); MOZ_RELEASE_ASSERT(iter.AdvanceAcrossSegments(bl2, kExtractSize)); MOZ_RELEASE_ASSERT(iter.Done()); + BufferList bl4(8, 8, 8); + bl4.WriteBytes("abcd1234", 8); + iter = bl4.Iter(); + iter.Advance(bl4, 8); + + BufferList bl5 = bl4.Extract(iter, kExtractSize, &success); + MOZ_RELEASE_ASSERT(!success); + + BufferList bl6(0, 0, 16); + bl6.WriteBytes("abcdefgh12345678", 16); + bl6.WriteBytes("ijklmnop87654321", 16); + iter = bl6.Iter(); + iter.Advance(bl6, 8); + BufferList bl7 = bl6.Extract(iter, 16, &success); + MOZ_RELEASE_ASSERT(success); + char data[16]; + MOZ_RELEASE_ASSERT(bl6.ReadBytes(iter, data, 8)); + MOZ_RELEASE_ASSERT(memcmp(data, "87654321", 8) == 0); + iter = bl7.Iter(); + MOZ_RELEASE_ASSERT(bl7.ReadBytes(iter, data, 16)); + MOZ_RELEASE_ASSERT(memcmp(data, "12345678ijklmnop", 16) == 0); + + BufferList bl8(0, 0, 16); + bl8.WriteBytes("abcdefgh12345678", 16); + iter = bl8.Iter(); + BufferList bl9 = bl8.Extract(iter, 8, &success); + MOZ_RELEASE_ASSERT(success); + MOZ_RELEASE_ASSERT(bl9.Size() == 8); + MOZ_RELEASE_ASSERT(!iter.Done()); + return 0; }