Read Ahead Bug Fixes

Author: Omar Marzouk

tensorflow_io/core/filesystems/dfs/dfs_filesystem.cc

@@ -52,25 +52,49 @@ void Cleanup(TF_RandomAccessFile* file) {
 int64_t Read(const TF_RandomAccessFile* file, uint64_t offset, size_t n,
              char* ret, TF_Status* status) {
   auto dfs_file = static_cast<DFSRandomAccessFile*>(file->plugin_file);
-  for (auto& read_buf : dfs_file->buffers) {
-    if (read_buf.CacheHit(offset, n))
-      return read_buf.CopyFromCache(ret, offset, n, dfs_file->file_size,
-                                    status);
-  }
+  if (offset > dfs_file->file_size) return -1;
+  size_t ret_offset = 0;
+  size_t curr_offset = offset;
+  int64_t total_bytes = 0;
+  size_t ret_size = offset + n;
+  while (curr_offset < ret_size && curr_offset < dfs_file->file_size) {
+    size_t read_bytes = 0;
+    for (auto& read_buf : dfs_file->buffers) {
+      if (read_buf.CacheHit(curr_offset)) {
+        read_bytes = read_buf.CopyFromCache(ret, ret_offset, offset, n,
+                                            dfs_file->file_size, status);
+      }
+    }
 
-  size_t curr_offset = offset + BUFF_SIZE;
-  for (size_t i = 1; i < dfs_file->buffers.size(); i++) {
-    if (curr_offset > dfs_file->file_size) break;
-    dfs_file->buffers[i].ReadAsync(dfs_file->daos_fs, dfs_file->daos_file.file,
-                                   curr_offset);
-    curr_offset += BUFF_SIZE;
-  }
+    if (read_bytes > 0) {
+      curr_offset += read_bytes;
+      ret_offset += read_bytes;
+      total_bytes += read_bytes;
+      n -= read_bytes;
+      continue;
+    }
 
-  dfs_file->buffers[0].ReadSync(dfs_file->daos_fs, dfs_file->daos_file.file,
-                                offset);
+    size_t async_offset = curr_offset + BUFF_SIZE;
+    for (size_t i = 1; i < dfs_file->buffers.size(); i++) {
+      if (async_offset > dfs_file->file_size) break;
+      dfs_file->buffers[i].ReadAsync(dfs_file->daos_fs,
+                                     dfs_file->daos_file.file, async_offset);
+      async_offset += BUFF_SIZE;
+    }
+
+    dfs_file->buffers[0].ReadSync(dfs_file->daos_fs, dfs_file->daos_file.file,
+                                  curr_offset);
+
+    read_bytes = dfs_file->buffers[0].CopyFromCache(
+        ret, ret_offset, curr_offset, n, dfs_file->file_size, status);
+
+    curr_offset += read_bytes;
+    ret_offset += read_bytes;
+    total_bytes += read_bytes;
+    n -= read_bytes;
+  }
 
-  return dfs_file->buffers[0].CopyFromCache(ret, offset, n, dfs_file->file_size,
-                                            status);
+  return total_bytes;
 }
 
 }  // namespace tf_random_access_file

tensorflow_io/core/filesystems/dfs/dfs_utils.cc

@@ -515,6 +515,7 @@ ReadBuffer::ReadBuffer(size_t id, daos_handle_t eqh, size_t size)
     : id(id), buffer_size(size), eqh(eqh) {
   buffer = new char[size];
   buffer_offset = 0;
+  initialized = false;
   event = new daos_event_t;
   daos_event_init(event, eqh, nullptr);
   valid = false;
@@ -540,11 +541,12 @@ ReadBuffer::ReadBuffer(ReadBuffer&& read_buffer) {
   valid = false;
   read_buffer.buffer = nullptr;
   read_buffer.event = nullptr;
+  initialized = false;
 }
 
-bool ReadBuffer::CacheHit(const size_t pos, const size_t len) {
-  return pos >= buffer_offset && len <= buffer_size &&
-         (pos + len <= buffer_offset + buffer_size);
+bool ReadBuffer::CacheHit(const size_t pos) {
+  return pos >= buffer_offset && (pos < buffer_offset + buffer_size) &&
+         initialized;
 }
 
 int ReadBuffer::WaitEvent() {
@@ -575,6 +577,7 @@ int ReadBuffer::ReadAsync(dfs_t* daos_fs, dfs_obj_t* file, const size_t off) {
   rsgl.sg_iovs = &iov;
   valid = false;
   buffer_offset = off;
+  initialized = true;
   dfs_read(daos_fs, file, &rsgl, buffer_offset, &read_size, event);
   return 0;
 }
@@ -587,23 +590,29 @@ int ReadBuffer::ReadSync(dfs_t* daos_fs, dfs_obj_t* file, const size_t off) {
   rsgl.sg_iovs = &iov;
   valid = false;
   buffer_offset = off;
+  initialized = true;
   rc = dfs_read(daos_fs, file, &rsgl, off, &read_size, NULL);
   if (!rc) valid = true;
   return rc;
 }
 
-int ReadBuffer::CopyData(char* ret, const size_t off, const size_t n) {
+int ReadBuffer::CopyData(char* ret, const size_t ret_offset, const size_t off,
+                         const size_t n) {
   int rc = WaitEvent();
   if (rc) return rc;
-  memcpy(ret, buffer + (off - buffer_offset), n);
+  memcpy(ret + ret_offset, buffer + (off - buffer_offset), n);
   return 0;
 }
 
-int ReadBuffer::CopyFromCache(char* ret, const size_t off, const size_t n,
+int ReadBuffer::CopyFromCache(char* ret, const size_t ret_offset,
+                              const size_t off, const size_t n,
                               const daos_size_t file_size, TF_Status* status) {
   size_t read_size;
   read_size = off + n > file_size ? file_size - off : n;
-  int rc = CopyData(ret, off, read_size);
+  read_size = off + read_size > buffer_offset + buffer_size
+                  ? buffer_offset + buffer_size - off
+                  : read_size;
+  int rc = CopyData(ret, ret_offset, off, read_size);
   if (rc) {
     TF_SetStatus(status, TF_INTERNAL, "");
     return 0;

tensorflow_io/core/filesystems/dfs/dfs_utils.h

@@ -215,7 +215,7 @@ class ReadBuffer {
 
   ~ReadBuffer();
 
-  bool CacheHit(const size_t pos, const size_t off);
+  bool CacheHit(const size_t pos);
 
   int WaitEvent();
 
@@ -225,10 +225,12 @@ class ReadBuffer {
 
   int ReadSync(dfs_t* dfs, dfs_obj_t* file, const size_t off);
 
-  int CopyData(char* ret, const size_t offset, const size_t n);
+  int CopyData(char* ret, const size_t ret_offset, const size_t offset,
+               const size_t n);
 
-  int CopyFromCache(char* ret, const size_t off, const size_t n,
-                    const daos_size_t file_size, TF_Status* status);
+  int CopyFromCache(char* ret, const size_t ret_offset, const size_t off,
+                    const size_t n, const daos_size_t file_size,
+                    TF_Status* status);
 
  private:
   size_t id;
@@ -241,6 +243,7 @@ class ReadBuffer {
   d_iov_t iov;
   bool valid;
   daos_size_t read_size;
+  bool initialized;
 };
 
 #endif  // TENSORFLOW_IO_CORE_FILESYSTEMS_DFS_DFS_FILESYSTEM_H_

tests/test_dfs.py

@@ -139,7 +139,7 @@ def test_list_directory(self):
         # Setup and check preconditions.
         dir_name = self._path_to("listdir")
         tf.io.gfile.mkdir(dir_name)
-        file_names = [self._path_to(f"listdir/{i}") for i in range(1, 4)]
+        file_names = [self._path_to(f"listdir/{i}") for i in range(1, 33)]
 
         for file_name in file_names:
             with tf.io.gfile.GFile(file_name, "w") as write_file: