dxva2n: link AVFrames and IMediaSample directly

This avoids issues with the external management getting confused on
(broken) flushes. This way at least playback remains stable.

3 files changed, 9 insertions, 21 deletions

diff --git a/decoder/LAVVideo/decoders/avcodec.cpp b/decoder/LAVVideo/decoders/avcodec.cpp
index 9b43ced5..764880ea 100644
--- a/decoder/LAVVideo/decoders/avcodec.cpp
+++ b/decoder/LAVVideo/decoders/avcodec.cpp
@@ -929,6 +929,7 @@ STDMETHODIMP CDecAvcodec::Decode(const BYTE *buffer, int buflen, REFERENCE_TIME
       pOutFrame->flags |= LAV_FRAME_FLAG_END_OF_SEQUENCE;
 
     if (pOutFrame->format == LAVPixFmt_DXVA2) {
+      pOutFrame->data[0] = m_pFrame->data[4];
       HandleDXVA2Frame(pOutFrame);
     } else {
       Deliver(pOutFrame);
diff --git a/decoder/LAVVideo/decoders/dxva2dec.cpp b/decoder/LAVVideo/decoders/dxva2dec.cpp
index e63421b8..940bc15f 100644
--- a/decoder/LAVVideo/decoders/dxva2dec.cpp
+++ b/decoder/LAVVideo/decoders/dxva2dec.cpp
@@ -198,10 +198,6 @@ STDMETHODIMP CDecDXVA2::DestroyDecoder(bool bFull, bool bNoAVCodec)
   m_pCallback->ReleaseAllDXVAResources();
   for (int i = 0; i < m_NumSurfaces; i++) {
     SafeRelease(&m_pSurfaces[i].d3d);
-    // To avoid a double free, remove the reference first, then call the release
-    IMediaSample *pSample = m_pSurfaces[i].sample;
-    m_pSurfaces[i].sample = NULL;
-    SafeRelease(&pSample);
   }
   m_NumSurfaces = 0;
 
@@ -902,7 +898,6 @@ HRESULT CDecDXVA2::CreateDXVA2Decoder(int nSurfaces, IDirect3DSurface9 **ppSurfa
     m_pSurfaces[i].index = i;
     m_pSurfaces[i].d3d = ppSurfaces[i];
     m_pSurfaces[i].age = 0;
-    m_pSurfaces[i].sample = NULL;
     m_pSurfaces[i].used = false;
   }
 
@@ -960,6 +955,7 @@ static enum PixelFormat get_dxva2_format(struct AVCodecContext *s, const enum Pi
 
 typedef struct SurfaceWrapper {
   LPDIRECT3DSURFACE9 surface;
+  IMediaSample *sample;
   CDecDXVA2 *pDec;
 } SurfaceWrapper;
 
@@ -971,19 +967,18 @@ void CDecDXVA2::free_dxva2_buffer(void *opaque, uint8_t *data)
   LPDIRECT3DSURFACE9 pSurface = sw->surface;
   for (int i = 0; i < pDec->m_NumSurfaces; i++) {
     if (pDec->m_pSurfaces[i].d3d == pSurface) {
-      IMediaSample *pSample = pDec->m_pSurfaces[i].sample;
-      pDec->m_pSurfaces[i].sample = NULL;
-      SafeRelease(&pSample);
       pDec->m_pSurfaces[i].used = false;
       break;
     }
   }
+  SafeRelease(&sw->sample);
   delete sw;
 }
 
 int CDecDXVA2::get_dxva2_buffer(struct AVCodecContext *c, AVFrame *pic, int flags)
 {
   CDecDXVA2 *pDec = (CDecDXVA2 *)c->opaque;
+  IMediaSample *pSample = NULL;
 
   HRESULT hr = S_OK;
 
@@ -1036,7 +1031,6 @@ int CDecDXVA2::get_dxva2_buffer(struct AVCodecContext *c, AVFrame *pic, int flag
     if (!pDec->m_pDXVA2Allocator)
       return -1;
 
-    IMediaSample *pSample = NULL;
     hr = pDec->m_pDXVA2Allocator->GetBuffer(&pSample, NULL, NULL, 0);
     if (FAILED(hr)) {
       DbgLog((LOG_ERROR, 10, L"DXVA2Allocator returned error, hr: 0x%x", hr));
@@ -1052,8 +1046,6 @@ int CDecDXVA2::get_dxva2_buffer(struct AVCodecContext *c, AVFrame *pic, int flag
     }
     i = pLavDXVA2->GetDXSurfaceId();
     SafeRelease(&pLavDXVA2);
-
-    pDec->m_pSurfaces[i].sample = pSample;
   } else {
     int old, old_unused;
     for (i = 0, old = 0, old_unused = -1; i < pDec->m_NumSurfaces; i++) {
@@ -1074,7 +1066,7 @@ int CDecDXVA2::get_dxva2_buffer(struct AVCodecContext *c, AVFrame *pic, int flag
   LPDIRECT3DSURFACE9 pSurface = pDec->m_pSurfaces[i].d3d;
   if (!pSurface) {
     DbgLog((LOG_ERROR, 10, L"There is a sample, but no D3D Surace? WTF?"));
-    SafeRelease(&pDec->m_pSurfaces[i].sample);
+    SafeRelease(&pSample);
     return -1;
   }
 
@@ -1086,10 +1078,12 @@ int CDecDXVA2::get_dxva2_buffer(struct AVCodecContext *c, AVFrame *pic, int flag
   memset(pic->buf, 0, sizeof(pic->buf));
 
   pic->data[0] = pic->data[3] = (uint8_t *)pSurface;
+  pic->data[4] = (uint8_t *)pSample;
 
   SurfaceWrapper *surfaceWrapper = new SurfaceWrapper();
   surfaceWrapper->pDec = pDec;
   surfaceWrapper->surface = pSurface;
+  surfaceWrapper->sample = pSample;
   pic->buf[0] = av_buffer_create(NULL, 0, free_dxva2_buffer, surfaceWrapper, 0);
 
   return 0;
@@ -1137,9 +1131,6 @@ STDMETHODIMP CDecDXVA2::Flush()
 
   for (int i = 0; i < m_NumSurfaces; i++) {
     d3d_surface_t *s = &m_pSurfaces[i];
-    IMediaSample *pSample = s->sample;
-    s->sample = NULL;
-    SafeRelease(&pSample);
     s->used = false;
     //s->age = 0;
   }
@@ -1204,15 +1195,12 @@ HRESULT CDecDXVA2::HandleDXVA2Frame(LAVFrame *pFrame)
 HRESULT CDecDXVA2::DeliverDXVA2Frame(LAVFrame *pFrame)
 {
   if (m_bNative) {
-    d3d_surface_t *s = FindSurface((LPDIRECT3DSURFACE9)pFrame->data[3]);
-
-    if (!s->sample) {
+    if (!pFrame->data[0] || !pFrame->data[3]) {
+      DbgLog((LOG_ERROR, 10, L"No sample or surface for DXVA2 frame?!?!"));
       ReleaseFrame(&pFrame);
       return S_FALSE;
     }
 
-    pFrame->data[0] = (uint8_t *)s->sample;
-
     pFrame->format = LAVPixFmt_DXVA2;
     Deliver(pFrame);
   } else {
diff --git a/decoder/LAVVideo/decoders/dxva2dec.h b/decoder/LAVVideo/decoders/dxva2dec.h
index 6273bb56..6071718f 100644
--- a/decoder/LAVVideo/decoders/dxva2dec.h
+++ b/decoder/LAVVideo/decoders/dxva2dec.h
@@ -31,7 +31,6 @@ typedef struct {
   bool used;
   LPDIRECT3DSURFACE9 d3d;
   uint64_t age;
-  IMediaSample *sample;
 } d3d_surface_t;
 
 class CDXVA2SurfaceAllocator;