瀏覽代碼

Fixed memory leak.

jcsyshc 2 年之前
父節點
當前提交
78240750a8
共有 4 個文件被更改,包括 21 次插入20 次删除
  1. 1 5
      CMakeLists.txt
  2. 9 5
      src/main.cpp
  3. 8 8
      src/video_encoder.cpp
  4. 3 2
      src/video_encoder.h

+ 1 - 5
CMakeLists.txt

@@ -117,8 +117,4 @@ endif ()
 set(NVCODEC_INCLUDE_DIR ${NVCODEC_DIR}/Interface)
 target_include_directories(${PROJECT_NAME} PRIVATE ${NVCODEC_INCLUDE_DIR})
 target_link_libraries(${PROJECT_NAME} ${NVENC_LIB})
-target_sources(${PROJECT_NAME} PRIVATE src/video_encoder.cpp)
-
-# WebRTC config
-find_package(LibDataChannel REQUIRED)
-target_link_libraries(${PROJECT_NAME} LibDataChannel::LibDataChannel)
+target_sources(${PROJECT_NAME} PRIVATE src/video_encoder.cpp)

+ 9 - 5
src/main.cpp

@@ -225,15 +225,19 @@ int main() {
 
             // encode frame
             output_fbo.download_pixels();
-            void *frame_data;
-            size_t frame_length;
-            encoder.encode_frame(output_fbo.pbo_res, &frame_data, &frame_length);
+            void *frame_ptr;
+            encoder.encode_frame(output_fbo.pbo_res, &frame_ptr);
+            auto frame_length = *(size_t *) ((char *) frame_ptr + 0);
+            auto frame_data = (char *) frame_ptr + sizeof(size_t);
 
-            SPDLOG_TRACE("Time used: {}ms", std::chrono::duration_cast<std::chrono::milliseconds>(
-                    std::chrono::high_resolution_clock::now() - start_time).count());
+            SPDLOG_TRACE("Time used: {}ms, length = {}", std::chrono::duration_cast<std::chrono::milliseconds>(
+                    std::chrono::high_resolution_clock::now() - start_time).count(), frame_length);
 
             // save encoded frame
             fwrite(frame_data, frame_length, 1, video_save_file);
+
+            // cleanup
+            free(frame_ptr);
         }
 
         int frame_width, frame_height;

+ 8 - 8
src/video_encoder.cpp

@@ -34,7 +34,6 @@ struct video_encoder::impl {
 
     // frame related
     void *frame_ptr = nullptr, **output_ptr = nullptr;
-    size_t *output_size = nullptr;
     NV_ENC_REGISTERED_PTR frame_reg_ptr = nullptr;
 
     bool initialize() {
@@ -168,9 +167,11 @@ struct video_encoder::impl {
         // copy bitstream
         assert(output_ptr != nullptr);
         assert(output_size != nullptr);
-        *output_ptr = malloc(lock_config.bitstreamSizeInBytes);
-        *output_size = lock_config.bitstreamSizeInBytes;
-        memcpy(*output_ptr, lock_config.bitstreamBufferPtr, lock_config.bitstreamSizeInBytes);
+        *output_ptr = malloc(lock_config.bitstreamSizeInBytes + sizeof(size_t));
+        auto &output_size = *(size_t *) ((char *) *output_ptr + 0);
+        auto output_data = (char *) *output_ptr + sizeof(size_t);
+        output_size = lock_config.bitstreamSizeInBytes;
+        memcpy(output_data, lock_config.bitstreamBufferPtr, lock_config.bitstreamSizeInBytes);
 
         // cleanup
         NVENC_API_CHECK(api.nvEncUnlockBitstream(encoder, output_buf));
@@ -205,7 +206,7 @@ void video_encoder::stop_encode() {
     pimpl->cleanup();
 }
 
-bool video_encoder::encode_frame(void *frame_ptr, void **output_ptr, size_t *output_size) {
+bool video_encoder::encode_frame(void *frame_ptr, void **output_ptr) {
     // register frame ptr
     if (pimpl->frame_ptr != frame_ptr) {
         pimpl->unregister_frame_ptr();
@@ -214,17 +215,16 @@ bool video_encoder::encode_frame(void *frame_ptr, void **output_ptr, size_t *out
     }
 
     pimpl->output_ptr = output_ptr;
-    pimpl->output_size = output_size;
     return pimpl->encode_frame();
 }
 
-bool video_encoder::encode_frame(cudaGraphicsResource *res, void **output_ptr, size_t *output_size) {
+bool video_encoder::encode_frame(cudaGraphicsResource *res, void **output_ptr) {
     void *pbo_ptr;
     size_t pbo_size;
     CUDA_API_CHECK(cudaGraphicsMapResources(1, &res));
     CUDA_API_CHECK(cudaGraphicsResourceGetMappedPointer(&pbo_ptr, &pbo_size, res));
     assert(pbo_size == pimpl->frame_pitch * pimpl->frame_height);
-    CALL_CHECK(encode_frame(pbo_ptr, output_ptr, output_size));
+    CALL_CHECK(encode_frame(pbo_ptr, output_ptr));
     CUDA_API_CHECK(cudaGraphicsUnmapResources(1, &res));
     return true;
 }

+ 3 - 2
src/video_encoder.h

@@ -18,9 +18,10 @@ public:
 
     void stop_encode();
 
-    bool encode_frame(void *frame_ptr, void **output_ptr, size_t *output_size);
+    // leading 4 bytes indicate length
+    bool encode_frame(void *frame_ptr, void **output_ptr);
 
-    bool encode_frame(cudaGraphicsResource *res, void **output_ptr, size_t *output_size);
+    bool encode_frame(cudaGraphicsResource *res, void **output_ptr);
 
     bool is_encoding();