Fix the memory leak in the video decoder

Limelight/Connection.m

@@ -65,9 +65,8 @@ void DrSubmitDecodeUnit(PDECODE_UNIT decodeUnit)
             entry = entry->next;
         }
         
+        // This function will take our buffer
         [renderer submitDecodeBuffer:data length:decodeUnit->fullLength];
-        
-        free(data);
     }
 }
 

Limelight/VideoDecoderRenderer.m

@@ -46,38 +46,69 @@
     OSStatus status;
     size_t oldOffset = CMBlockBufferGetDataLength(existingBuffer);
     
-    // Append a 4 byte buffer to this block for the length prefix
+    // If we're at index 1 (first NALU in frame), enqueue this buffer to the memory block
-    status = CMBlockBufferAppendMemoryBlock(existingBuffer, NULL,
+    // so it can handle freeing it when the block buffer is destroyed
-                                            NAL_LENGTH_PREFIX_SIZE,
+    if (offset == 1) {
-                                            kCFAllocatorDefault, NULL, 0,
+        int dataLength = nalLength - NALU_START_PREFIX_SIZE;
-                                            NAL_LENGTH_PREFIX_SIZE, 0);
-    if (status != noErr) {
-        NSLog(@"CMBlockBufferAppendMemoryBlock failed: %d", (int)status);
-        return;
-    }
         
-    // Write the length prefix to the new buffer
+        // Pass the real buffer pointer directly (no offset)
-    int dataLength = nalLength - NALU_START_PREFIX_SIZE;
+        // This will give it to the block buffer to free when it's released.
-    const uint8_t lengthBytes[] = {(uint8_t)(dataLength >> 24), (uint8_t)(dataLength >> 16),
+        // All further calls to CMBlockBufferAppendMemoryBlock will do so
-        (uint8_t)(dataLength >> 8), (uint8_t)dataLength};
+        // at an offset and will not be asking the buffer to be freed.
-    status = CMBlockBufferReplaceDataBytes(lengthBytes, existingBuffer,
+        status = CMBlockBufferAppendMemoryBlock(existingBuffer, data,
-                                           oldOffset, NAL_LENGTH_PREFIX_SIZE);
+                                                nalLength + 1, // Add 1 for the offset we decremented
-    if (status != noErr) {
+                                                kCFAllocatorDefault,
-        NSLog(@"CMBlockBufferReplaceDataBytes failed: %d", (int)status);
+                                                NULL, 0, nalLength + 1, 0);
-        return;
+        if (status != noErr) {
-    }
+            NSLog(@"CMBlockBufferReplaceDataBytes failed: %d", (int)status);
+            return;
+        }
         
-    // Append the rest of the data by simply attaching a reference to the buffer
+        // Write the length prefix to existing buffer
-    status = CMBlockBufferAppendMemoryBlock(existingBuffer, &data[offset+NALU_START_PREFIX_SIZE],
+        const uint8_t lengthBytes[] = {(uint8_t)(dataLength >> 24), (uint8_t)(dataLength >> 16),
-                                            dataLength,
+            (uint8_t)(dataLength >> 8), (uint8_t)dataLength};
-                                            kCFAllocatorNull, // Don't deallocate data on free
+        status = CMBlockBufferReplaceDataBytes(lengthBytes, existingBuffer,
-                                            NULL, 0, dataLength, 0);
+                                               oldOffset, NAL_LENGTH_PREFIX_SIZE);
-    if (status != noErr) {
+        if (status != noErr) {
-        NSLog(@"CMBlockBufferReplaceDataBytes failed: %d", (int)status);
+            NSLog(@"CMBlockBufferReplaceDataBytes failed: %d", (int)status);
-        return;
+            return;
+        }
+    }
+    else {
+        // Append a 4 byte buffer to this block for the length prefix
+        status = CMBlockBufferAppendMemoryBlock(existingBuffer, NULL,
+                                                NAL_LENGTH_PREFIX_SIZE,
+                                                kCFAllocatorDefault, NULL, 0,
+                                                NAL_LENGTH_PREFIX_SIZE, 0);
+        if (status != noErr) {
+            NSLog(@"CMBlockBufferAppendMemoryBlock failed: %d", (int)status);
+            return;
+        }
+        
+        // Write the length prefix to the new buffer
+        int dataLength = nalLength - NALU_START_PREFIX_SIZE;
+        const uint8_t lengthBytes[] = {(uint8_t)(dataLength >> 24), (uint8_t)(dataLength >> 16),
+            (uint8_t)(dataLength >> 8), (uint8_t)dataLength};
+        status = CMBlockBufferReplaceDataBytes(lengthBytes, existingBuffer,
+                                               oldOffset, NAL_LENGTH_PREFIX_SIZE);
+        if (status != noErr) {
+            NSLog(@"CMBlockBufferReplaceDataBytes failed: %d", (int)status);
+            return;
+        }
+        
+        // Attach the buffer by reference to the block buffer
+        status = CMBlockBufferAppendMemoryBlock(existingBuffer, &data[offset+NALU_START_PREFIX_SIZE],
+                                                dataLength,
+                                                kCFAllocatorNull, // Don't deallocate data on free
+                                                NULL, 0, dataLength, 0);
+        if (status != noErr) {
+            NSLog(@"CMBlockBufferReplaceDataBytes failed: %d", (int)status);
+            return;
+        }
     }
 }
 
+// This function must free data
 - (void)submitDecodeBuffer:(unsigned char *)data length:(int)length
 {
     unsigned char nalType = data[FRAME_START_PREFIX_SIZE] & 0x1F;
@@ -120,21 +151,25 @@
             if (status != noErr) {
                 NSLog(@"Failed to create format description: %d", (int)status);
                 formatDesc = NULL;
-                return;
             }
         }
         
+        // Free the data buffer
+        free(data);
+        
         // No frame data to submit for these NALUs
         return;
     }
     
     if (formatDesc == NULL) {
         // Can't decode if we haven't gotten our parameter sets yet
+        free(data);
         return;
     }
     
     if (nalType != 0x1 && nalType != 0x5) {
         // Don't submit parameter set data
+        free(data);
         return;
     }
     
@@ -144,6 +179,7 @@
     status = CMBlockBufferCreateEmpty(NULL, 0, 0, &blockBuffer);
     if (status != noErr) {
         NSLog(@"CMBlockBufferCreateEmpty failed: %d", (int)status);
+        free(data);
         return;
     }
     
@@ -166,6 +202,8 @@
         [self updateBufferForRange:blockBuffer data:data offset:lastOffset length:length - lastOffset];
     }
     
+    // From now on, CMBlockBuffer owns the data pointer and will free it when it's dereferenced
+    
     CMSampleBufferRef sampleBuffer;
     
     status = CMSampleBufferCreate(kCFAllocatorDefault,
@@ -176,6 +214,7 @@
                                   &sampleBuffer);
     if (status != noErr) {
         NSLog(@"CMSampleBufferCreate failed: %d", (int)status);
+        CFRelease(blockBuffer);
         return;
     }
     
@@ -197,6 +236,10 @@
     }
     
     [displayLayer enqueueSampleBuffer:sampleBuffer];
+    
+    // Dereference the buffers
+    CFRelease(blockBuffer);
+    CFRelease(sampleBuffer);
 }
 
 @end