From d887413eca3b683c2f3d110c01e7e680c7072548 Mon Sep 17 00:00:00 2001
From: Angie Chiang <angiebird@google.com>
Date: Thu, 14 Mar 2019 12:28:25 -0700
Subject: [PATCH] Fix race condition in wiener_var_rd_mult

Change-Id: Id5e9c2cbfe35809ac99a3bc9ba93cf462a6b1a34
---
 vp9/encoder/vp9_encodeframe.c | 6 ++++++
 vp9/encoder/vp9_encoder.c     | 6 ++++++
 vp9/encoder/vp9_encoder.h     | 3 +++
 3 files changed, 15 insertions(+)

diff --git a/vp9/encoder/vp9_encodeframe.c b/vp9/encoder/vp9_encodeframe.c
index a0e531e3d..047b06b53 100644
--- a/vp9/encoder/vp9_encodeframe.c
+++ b/vp9/encoder/vp9_encodeframe.c
@@ -3595,9 +3595,15 @@ static int wiener_var_rdmult(VP9_COMP *cpi, BLOCK_SIZE bsize, int mi_row,
     for (col = mb_col_start; col < mb_col_end; ++col)
       wiener_variance += cpi->mb_wiener_variance[row * cm->mb_cols + col];
 
+#if CONFIG_MULTITHREAD
+  pthread_mutex_lock(&cpi->kmeans_mutex);
+#endif  // CONFIG_MULTITHREAD
   kmeans_data = &cpi->kmeans_data_arr[cpi->kmeans_data_size++];
   kmeans_data->value = log(1 + wiener_variance);
   kmeans_data->pos = mi_row * cpi->kmeans_data_stride + mi_col;
+#if CONFIG_MULTITHREAD
+  pthread_mutex_unlock(&cpi->kmeans_mutex);
+#endif  // CONFIG_MULTITHREAD
   if (wiener_variance)
     wiener_variance /=
         (mb_row_end - mb_row_start) * (mb_col_end - mb_col_start);
diff --git a/vp9/encoder/vp9_encoder.c b/vp9/encoder/vp9_encoder.c
index ccaa0815d..2a9a36f61 100644
--- a/vp9/encoder/vp9_encoder.c
+++ b/vp9/encoder/vp9_encoder.c
@@ -2593,6 +2593,9 @@ void vp9_remove_compressor(VP9_COMP *cpi) {
 #endif
 
   if (cpi->kmeans_data_arr_alloc) {
+#if CONFIG_MULTITHREAD
+    pthread_mutex_destroy(&cpi->kmeans_mutex);
+#endif
     vpx_free(cpi->kmeans_data_arr);
   }
 
@@ -7263,6 +7266,9 @@ int vp9_get_compressed_data(VP9_COMP *cpi, unsigned int *frame_flags,
   if (cpi->kmeans_data_arr_alloc == 0) {
     const int mi_cols = mi_cols_aligned_to_sb(cm->mi_cols);
     const int mi_rows = mi_cols_aligned_to_sb(cm->mi_rows);
+#if CONFIG_MULTITHREAD
+    pthread_mutex_init(&cpi->kmeans_mutex, NULL);
+#endif
     CHECK_MEM_ERROR(
         cm, cpi->kmeans_data_arr,
         vpx_calloc(mi_rows * mi_cols, sizeof(*cpi->kmeans_data_arr)));
diff --git a/vp9/encoder/vp9_encoder.h b/vp9/encoder/vp9_encoder.h
index 24ccaf99d..567f05934 100644
--- a/vp9/encoder/vp9_encoder.h
+++ b/vp9/encoder/vp9_encoder.h
@@ -602,6 +602,9 @@ typedef struct VP9_COMP {
   TplDepFrame tpl_stats[MAX_ARF_GOP_SIZE];
   YV12_BUFFER_CONFIG *tpl_recon_frames[REF_FRAMES];
   EncFrameBuf enc_frame_buf[REF_FRAMES];
+#if CONFIG_MULTITHREAD
+  pthread_mutex_t kmeans_mutex;
+#endif
   int kmeans_data_arr_alloc;
   KMEANS_DATA *kmeans_data_arr;
   int kmeans_data_size;
-- 
2.40.0