Add memory fences to ensure that the reader's deliberate racy use of

is_initialized is safe.
Change: 118405243

tensorflow/core/kernels/initializable_lookup_table.cc

@@ -25,6 +25,9 @@ Status InitializableLookupTable::Find(const Tensor& keys, Tensor* values,
   if (!is_initialized()) {
     return errors::FailedPrecondition("Table not initialized.");
   }
+  // Do not let the use migrate before the check;  table is used without
+  // a lock by the readers.
+  std::atomic_thread_fence(std::memory_order_acquire);
   TF_RETURN_IF_ERROR(CheckFindArguments(keys, *values, default_value));
   return DoFind(keys, values, default_value);
 }
@@ -48,6 +51,10 @@ Status InitializableLookupTable::Initialize(InitTableIterator& iter) {
   if (!errors::IsOutOfRange(iter.status())) {
     return iter.status();
   }
+
+  // Prevent compiler/memory reordering of is_initialized and
+  // the initialization itself.
+  std::atomic_thread_fence(std::memory_order_release);
   is_initialized_ = true;
   return Status::OK();
 }

tensorflow/core/kernels/lookup_table_op.cc

@@ -69,7 +69,11 @@ class HashTable : public InitializableLookupTable {
  public:
   size_t size() const override {
     // return the size of the table only if it's initialized, otherwise 0.
-    return table_ && is_initialized_ ? table_->size() : 0;
+    if (!is_initialized_) {
+      return 0;
+    }
+    std::atomic_thread_fence(std::memory_order_acquire);
+    return table_ ? table_->size() : 0;
   }
 
   DataType key_dtype() const override { return DataTypeToEnum<K>::v(); }