@@ -370,8 +370,68 @@ class ZstdDecompressContext final : public ZstdContext { |
370 | 370 | DeleteFnPtr<ZSTD_DCtx, ZstdDecompressContext::FreeZstd> dctx_; |
371 | 371 | }; |
372 | 372 |
|
| 373 | +class CompressionStreamMemoryOwner { |
| 374 | + public: |
| 375 | + // Allocation functions provided to zlib itself. We store the real size of |
| 376 | + // the allocated memory chunk just before the "payload" memory we return |
| 377 | + // to zlib. |
| 378 | + // Because we use zlib off the thread pool, we can not report memory directly |
| 379 | + // to V8; rather, we first store it as "unreported" memory in a separate |
| 380 | + // field and later report it back from the main thread. |
| 381 | + static void* AllocForZlib(void* data, uInt items, uInt size) { |
| 382 | + size_t real_size = MultiplyWithOverflowCheck(static_cast<size_t>(items), |
| 383 | + static_cast<size_t>(size)); |
| 384 | + return AllocForBrotli(data, real_size); |
| 385 | + } |
| 386 | + |
| 387 | + static constexpr size_t reserveSizeAndAlign = |
| 388 | + std::max(sizeof(size_t), alignof(max_align_t)); |
| 389 | + |
| 390 | + static void* AllocForBrotli(void* data, size_t size) { |
| 391 | + size += reserveSizeAndAlign; |
| 392 | + CompressionStreamMemoryOwner* ctx = |
| 393 | + static_cast<CompressionStreamMemoryOwner*>(data); |
| 394 | + char* memory = UncheckedMalloc(size); |
| 395 | + if (memory == nullptr) [[unlikely]] { |
| 396 | + return nullptr; |
| 397 | + } |
| 398 | + *reinterpret_cast<size_t*>(memory) = size; |
| 399 | + ctx->unreported_allocations_.fetch_add(size, std::memory_order_relaxed); |
| 400 | + return memory + reserveSizeAndAlign; |
| 401 | + } |
| 402 | + |
| 403 | + static void FreeForZlib(void* data, void* pointer) { |
| 404 | + if (pointer == nullptr) [[unlikely]] { |
| 405 | + return; |
| 406 | + } |
| 407 | + CompressionStreamMemoryOwner* ctx = |
| 408 | + static_cast<CompressionStreamMemoryOwner*>(data); |
| 409 | + char* real_pointer = static_cast<char*>(pointer) - reserveSizeAndAlign; |
| 410 | + size_t real_size = *reinterpret_cast<size_t*>(real_pointer); |
| 411 | + ctx->unreported_allocations_.fetch_sub(real_size, |
| 412 | + std::memory_order_relaxed); |
| 413 | + free(real_pointer); |
| 414 | + } |
| 415 | + |
| 416 | + void* as_allocator_opaque_value() { return static_cast<void*>(this); } |
| 417 | + |
| 418 | + protected: |
| 419 | + ssize_t ComputeAdjustmentToExternalAllocatedMemory() { |
| 420 | + ssize_t report = |
| 421 | + unreported_allocations_.exchange(0, std::memory_order_relaxed); |
| 422 | + CHECK_IMPLIES(report < 0, zlib_memory_ >= static_cast<size_t>(-report)); |
| 423 | + zlib_memory_ += report; |
| 424 | + return report; |
| 425 | + } |
| 426 | + |
| 427 | + std::atomic<ssize_t> unreported_allocations_{0}; |
| 428 | + size_t zlib_memory_ = 0; |
| 429 | +}; |
| 430 | + |
373 | 431 | template <typename CompressionContext> |
374 | | -class CompressionStream : public AsyncWrap, public ThreadPoolWork { |
| 432 | +class CompressionStream : public AsyncWrap, |
| 433 | + public ThreadPoolWork, |
| 434 | + protected CompressionStreamMemoryOwner { |
375 | 435 | public: |
376 | 436 | enum InternalFields { |
377 | 437 | kCompressionStreamBaseField = AsyncWrap::kInternalFieldCount, |
@@ -591,6 +651,19 @@ class CompressionStream : public AsyncWrap, public ThreadPoolWork { |
591 | 651 | zlib_memory_ + unreported_allocations_); |
592 | 652 | } |
593 | 653 |
|
| 654 | + static void* AllocatorOpaquePointerForContext(CompressionContext* ctx) { |
| 655 | + CompressionStream* self = ContainerOf(&CompressionStream::ctx_, ctx); |
| 656 | + // There is nothing at the type level stopping someone from using this |
| 657 | + // method with `ctx` being an argument that is not part of a |
| 658 | + // CompressionStream. This check catches that potential discrepancy in debug |
| 659 | + // builds. std::launder is necessary to keep the compiler from optimizing |
| 660 | + // away the check in the (common) case that `CompressionContext` is a final |
| 661 | + // class. |
| 662 | + DCHECK_EQ(std::launder<MemoryRetainer>(&self->ctx_)->MemoryInfoName(), |
| 663 | + CompressionContext{}.MemoryInfoName()); |
| 664 | + return self->as_allocator_opaque_value(); |
| 665 | + } |
| 666 | + |
594 | 667 | protected: |
595 | 668 | CompressionContext* context() { return &ctx_; } |
596 | 669 |
|
@@ -600,55 +673,11 @@ class CompressionStream : public AsyncWrap, public ThreadPoolWork { |
600 | 673 | init_done_ = true; |
601 | 674 | } |
602 | 675 |
|
603 | | - // Allocation functions provided to zlib itself. We store the real size of |
604 | | - // the allocated memory chunk just before the "payload" memory we return |
605 | | - // to zlib. |
606 | | - // Because we use zlib off the thread pool, we can not report memory directly |
607 | | - // to V8; rather, we first store it as "unreported" memory in a separate |
608 | | - // field and later report it back from the main thread. |
609 | | - static void* AllocForZlib(void* data, uInt items, uInt size) { |
610 | | - size_t real_size = |
611 | | - MultiplyWithOverflowCheck(static_cast<size_t>(items), |
612 | | - static_cast<size_t>(size)); |
613 | | - return AllocForBrotli(data, real_size); |
614 | | - } |
615 | | - |
616 | | - static constexpr size_t reserveSizeAndAlign = |
617 | | - std::max(sizeof(size_t), alignof(max_align_t)); |
618 | | - |
619 | | - static void* AllocForBrotli(void* data, size_t size) { |
620 | | - size += reserveSizeAndAlign; |
621 | | - CompressionStream* ctx = static_cast<CompressionStream*>(data); |
622 | | - char* memory = UncheckedMalloc(size); |
623 | | - if (memory == nullptr) [[unlikely]] { |
624 | | - return nullptr; |
625 | | - } |
626 | | - *reinterpret_cast<size_t*>(memory) = size; |
627 | | - ctx->unreported_allocations_.fetch_add(size, |
628 | | - std::memory_order_relaxed); |
629 | | - return memory + reserveSizeAndAlign; |
630 | | - } |
631 | | - |
632 | | - static void FreeForZlib(void* data, void* pointer) { |
633 | | - if (pointer == nullptr) [[unlikely]] { |
634 | | - return; |
635 | | - } |
636 | | - CompressionStream* ctx = static_cast<CompressionStream*>(data); |
637 | | - char* real_pointer = static_cast<char*>(pointer) - reserveSizeAndAlign; |
638 | | - size_t real_size = *reinterpret_cast<size_t*>(real_pointer); |
639 | | - ctx->unreported_allocations_.fetch_sub(real_size, |
640 | | - std::memory_order_relaxed); |
641 | | - free(real_pointer); |
642 | | - } |
643 | | - |
644 | 676 | // This is called on the main thread after zlib may have allocated something |
645 | 677 | // in order to report it back to V8. |
646 | 678 | void AdjustAmountOfExternalAllocatedMemory() { |
647 | | - ssize_t report = |
648 | | - unreported_allocations_.exchange(0, std::memory_order_relaxed); |
| 679 | + ssize_t report = ComputeAdjustmentToExternalAllocatedMemory(); |
649 | 680 | if (report == 0) return; |
650 | | - CHECK_IMPLIES(report < 0, zlib_memory_ >= static_cast<size_t>(-report)); |
651 | | - zlib_memory_ += report; |
652 | 681 | AsyncWrap::env()->external_memory_accounter()->Update( |
653 | 682 | AsyncWrap::env()->isolate(), report); |
654 | 683 | } |
@@ -679,8 +708,6 @@ class CompressionStream : public AsyncWrap, public ThreadPoolWork { |
679 | 708 | bool closed_ = false; |
680 | 709 | unsigned int refs_ = 0; |
681 | 710 | uint32_t* write_result_ = nullptr; |
682 | | - std::atomic<ssize_t> unreported_allocations_{0}; |
683 | | - size_t zlib_memory_ = 0; |
684 | 711 |
|
685 | 712 | CompressionContext ctx_; |
686 | 713 | }; |
@@ -757,7 +784,7 @@ class ZlibStream final : public CompressionStream<ZlibContext> { |
757 | 784 |
|
758 | 785 | AllocScope alloc_scope(wrap); |
759 | 786 | wrap->context()->SetAllocationFunctions( |
760 | | - AllocForZlib, FreeForZlib, static_cast<CompressionStream*>(wrap)); |
| 787 | + AllocForZlib, FreeForZlib, wrap->as_allocator_opaque_value()); |
761 | 788 | wrap->context()->Init(level, window_bits, mem_level, strategy, |
762 | 789 | std::move(dictionary)); |
763 | 790 | } |
@@ -819,11 +846,10 @@ class BrotliCompressionStream final : |
819 | 846 | wrap->InitStream(write_result, write_js_callback); |
820 | 847 |
|
821 | 848 | AllocScope alloc_scope(wrap); |
822 | | - CompressionError err = |
823 | | - wrap->context()->Init( |
824 | | - CompressionStream<CompressionContext>::AllocForBrotli, |
825 | | - CompressionStream<CompressionContext>::FreeForZlib, |
826 | | - static_cast<CompressionStream<CompressionContext>*>(wrap)); |
| 849 | + CompressionError err = wrap->context()->Init( |
| 850 | + CompressionStream<CompressionContext>::AllocForBrotli, |
| 851 | + CompressionStream<CompressionContext>::FreeForZlib, |
| 852 | + wrap->as_allocator_opaque_value()); |
827 | 853 | if (err.IsError()) { |
828 | 854 | wrap->EmitError(err); |
829 | 855 | // TODO(addaleax): Sometimes we generate better error codes in C++ land, |
|