Home | History | Annotate | Download | only in default
      1 // Copyright 2017 The TensorFlow Authors. All Rights Reserved.
      2 //
      3 // Licensed under the Apache License, Version 2.0 (the "License");
      4 // you may not use this file except in compliance with the License.
      5 // You may obtain a copy of the License at
      6 //
      7 //     http://www.apache.org/licenses/LICENSE-2.0
      8 //
      9 // Unless required by applicable law or agreed to in writing, software
     10 // distributed under the License is distributed on an "AS IS" BASIS,
     11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
     12 // See the License for the specific language governing permissions and
     13 // limitations under the License.
     14 // =============================================================================
     15 
     16 #include "tensorflow/contrib/ffmpeg/ffmpeg_lib.h"
     17 
     18 #include <array>
     19 #include <set>
     20 #include <string>
     21 #include <vector>
     22 
     23 #include "tensorflow/core/lib/core/threadpool.h"
     24 #include "tensorflow/core/lib/io/path.h"
     25 #include "tensorflow/core/platform/env.h"
     26 #include "tensorflow/core/platform/mutex.h"
     27 #include "tensorflow/core/platform/test.h"
     28 
     29 namespace tensorflow {
     30 namespace ffmpeg {
     31 namespace {
     32 
     33 TEST(FfmpegLibTest, TestTempDirectoryThreading) {
     34   // Testing a fix for a bug that allowed different threads to create
     35   // conflicting temp files.
     36   // See github.com/tensorflow/tensorflow/issues/5804 for details.
     37   const int32 kNumThreads = 10;
     38   const int32 kNumWorkItems = 10000;
     39   static constexpr size_t kStringsPerItem = 100;
     40   Env* environment = Env::Default();
     41   thread::ThreadPool pool(environment, "test", kNumThreads);
     42 
     43   mutex mu;
     44   std::vector<string> temp_filenames;
     45   temp_filenames.reserve(kNumWorkItems * kStringsPerItem);
     46 
     47   // Queue a large number of work items for the threads to process. Each work
     48   // item creates a temp file and then deletes it.
     49   for (int i = 0; i < kNumWorkItems; ++i) {
     50     pool.Schedule([&mu, &temp_filenames, environment]() {
     51       std::array<string, kStringsPerItem> buffer;
     52       for (int32 j = 0; j < kStringsPerItem; ++j) {
     53         buffer[j] = io::GetTempFilename("mp3");
     54         TF_QCHECK_OK(environment->DeleteFile(buffer[j]));
     55       }
     56       mutex_lock l(mu);
     57       for (const auto& fn : buffer) {
     58         temp_filenames.push_back(fn);
     59       }
     60     });
     61   }
     62 
     63   // Wait until all work items are complete.
     64   while (true) {
     65     mutex_lock l(mu);
     66     if (temp_filenames.size() == kNumWorkItems * kStringsPerItem) {
     67       break;
     68     }
     69   }
     70 
     71   // Check that no duplicates are created.
     72   std::set<string> unique_filenames;
     73   mutex_lock l(mu);
     74   for (const auto& fn : temp_filenames) {
     75     ASSERT_TRUE(unique_filenames.insert(fn).second);
     76   }
     77 }
     78 
     79 }  // namespace
     80 }  // namespace ffmpeg
     81 }  // namespace tensorflow
     82