[IO] Return bytes written in Stream::Write (#686)

Lunderberg · web-flow · commit 3031e4a61a98 · 2024-05-21T18:36:18.000-07:00
This commit updates the `Stream::Write` method to return the number of
bytes written, analogous to the current behavior of `Stream::Read`.
This information is necessary to correctly model partial writes to a
buffered stream.  (e.g. Writing to a OS pipe when the pipe's buffer is
smaller than the data to be written.)
diff --git a/include/dmlc/io.h b/include/dmlc/io.h
@@ -32,16 +32,17 @@ class Stream {  // NOLINT(*)
   /*!
    * \brief reads data from a stream
    * \param ptr pointer to a memory buffer
-   * \param size block size
-   * \return the size of data read
+   * \param size The maximum number of bytes to read
+   * \return The number of bytes read from the stream
    */
-  virtual size_t Read(void *ptr, size_t size) = 0;
+  virtual size_t Read(void* ptr, size_t size) = 0;
   /*!
    * \brief writes data to a stream
    * \param ptr pointer to a memory buffer
-   * \param size block size
+   * \param size The maximum number of bytes to write
+   * \return The number of bytes written
    */
-  virtual void Write(const void *ptr, size_t size) = 0;
+  virtual size_t Write(const void* ptr, size_t size) = 0;
   /*! \brief virtual destructor */
   virtual ~Stream(void) {}
   /*!
diff --git a/include/dmlc/memory_io.h b/include/dmlc/memory_io.h
@@ -30,23 +30,24 @@ struct MemoryFixedSizeStream : public SeekStream {
         buffer_size_(buffer_size) {
     curr_ptr_ = 0;
   }
-  virtual size_t Read(void *ptr, size_t size) {
+  virtual size_t Read(void *ptr, size_t size) override {
     CHECK(curr_ptr_ + size <= buffer_size_);
     size_t nread = std::min(buffer_size_ - curr_ptr_, size);
     if (nread != 0) std::memcpy(ptr, p_buffer_ + curr_ptr_, nread);
     curr_ptr_ += nread;
     return nread;
   }
-  virtual void Write(const void *ptr, size_t size) {
-    if (size == 0) return;
+  virtual size_t Write(const void *ptr, size_t size) override {
+    if (size == 0) return 0;
     CHECK(curr_ptr_ + size <=  buffer_size_);
     std::memcpy(p_buffer_ + curr_ptr_, ptr, size);
     curr_ptr_ += size;
+    return size;
   }
-  virtual void Seek(size_t pos) {
+  virtual void Seek(size_t pos) override {
     curr_ptr_ = static_cast<size_t>(pos);
   }
-  virtual size_t Tell(void) {
+  virtual size_t Tell(void) override {
     return curr_ptr_;
   }
 
@@ -73,25 +74,26 @@ struct MemoryStringStream : public dmlc::SeekStream {
       : p_buffer_(p_buffer) {
     curr_ptr_ = 0;
   }
-  virtual size_t Read(void *ptr, size_t size) {
+  virtual size_t Read(void *ptr, size_t size) override {
     CHECK(curr_ptr_ <= p_buffer_->length());
     size_t nread = std::min(p_buffer_->length() - curr_ptr_, size);
     if (nread != 0) std::memcpy(ptr, &(*p_buffer_)[0] + curr_ptr_, nread);
     curr_ptr_ += nread;
     return nread;
   }
-  virtual void Write(const void *ptr, size_t size) {
-    if (size == 0) return;
+  virtual size_t Write(const void *ptr, size_t size) override {
+    if (size == 0) return 0;
     if (curr_ptr_ + size > p_buffer_->length()) {
       p_buffer_->resize(curr_ptr_+size);
     }
     std::memcpy(&(*p_buffer_)[0] + curr_ptr_, ptr, size);
     curr_ptr_ += size;
+    return size;
   }
-  virtual void Seek(size_t pos) {
+  virtual void Seek(size_t pos) override {
     curr_ptr_ = static_cast<size_t>(pos);
   }
-  virtual size_t Tell(void) {
+  virtual size_t Tell(void) override {
     return curr_ptr_;
   }
 
diff --git a/src/io/hdfs_filesys.cc b/src/io/hdfs_filesys.cc
@@ -28,7 +28,7 @@ class HDFSStream : public SeekStream {
     }
   }
 
-  virtual size_t Read(void *ptr, size_t size) {
+  virtual size_t Read(void *ptr, size_t size) override {
     char *buf = static_cast<char*>(ptr);
     size_t nleft = size;
     size_t nmax = static_cast<size_t>(std::numeric_limits<tSize>::max());
@@ -48,7 +48,7 @@ class HDFSStream : public SeekStream {
     return size - nleft;
   }
 
-  virtual void Write(const void *ptr, size_t size) {
+  virtual size_t Write(const void *ptr, size_t size) override {
     const char *buf = reinterpret_cast<const char*>(ptr);
     size_t nleft = size;
     // When using builtin-java classes to write, the maximum write size
@@ -70,14 +70,15 @@ class HDFSStream : public SeekStream {
         LOG(FATAL) << "HDFSStream.hdfsWrite Error:" << strerror(errsv);
       }
     }
+    return size - nleft;
   }
-  virtual void Seek(size_t pos) {
+  virtual void Seek(size_t pos) override {
     if (hdfsSeek(fs_, fp_, pos) != 0) {
       int errsv = errno;
       LOG(FATAL) << "HDFSStream.hdfsSeek Error:" << strerror(errsv);
     }
   }
-  virtual size_t Tell(void) {
+  virtual size_t Tell(void) override {
     tOffset offset = hdfsTell(fs_, fp_);
     if (offset == -1) {
       int errsv = errno;
diff --git a/src/io/local_filesys.cc b/src/io/local_filesys.cc
@@ -31,21 +31,22 @@ class FileStream : public SeekStream {
   virtual ~FileStream(void) {
     this->Close();
   }
-  virtual size_t Read(void *ptr, size_t size) {
+  virtual size_t Read(void *ptr, size_t size) override {
     return std::fread(ptr, 1, size, fp_);
   }
-  virtual void Write(const void *ptr, size_t size) {
+  virtual size_t Write(const void *ptr, size_t size) override {
     CHECK(std::fwrite(ptr, 1, size, fp_) == size)
         << "FileStream.Write incomplete";
+    return 0;
   }
-  virtual void Seek(size_t pos) {
+  virtual void Seek(size_t pos) override {
 #ifndef _MSC_VER
     CHECK(!std::fseek(fp_, static_cast<long>(pos), SEEK_SET));  // NOLINT(*)
 #else  // _MSC_VER
     CHECK(!_fseeki64(fp_, pos, SEEK_SET));
 #endif  // _MSC_VER
   }
-  virtual size_t Tell(void) {
+  virtual size_t Tell(void) override {
 #ifndef _MSC_VER
     return std::ftell(fp_);
 #else  // _MSC_VER
diff --git a/src/io/s3_filesys.cc b/src/io/s3_filesys.cc
@@ -424,23 +424,24 @@ class CURLReadStreamBase : public SeekStream {
   virtual ~CURLReadStreamBase() {
     this->Cleanup();
   }
-  virtual size_t Tell(void) {
+  virtual size_t Tell(void) override {
     return curr_bytes_;
   }
   virtual bool AtEnd(void) const {
     return at_end_;
   }
-  virtual void Write(const void *ptr, size_t size) {
+  virtual size_t Write(const void *ptr, size_t size) override {
     LOG(FATAL) << "CURL.ReadStream cannot be used for write";
+    return 0;
   }
   // lazy seek function
-  virtual void Seek(size_t pos) {
+  virtual void Seek(size_t pos) override {
     if (curr_bytes_ != pos) {
       this->Cleanup();
       curr_bytes_ = pos;
     }
   }
-  virtual size_t Read(void *ptr, size_t size);
+  virtual size_t Read(void *ptr, size_t size) override ;
 
  protected:
   CURLReadStreamBase()
@@ -790,11 +791,11 @@ class WriteStream : public Stream {
     ecurl_ = curl_easy_init();
     this->Init();
   }
-  virtual size_t Read(void *ptr, size_t size) {
+  virtual size_t Read(void *ptr, size_t size) override {
     LOG(FATAL) << "S3.WriteStream cannot be used for read";
     return 0;
   }
-  virtual void Write(const void *ptr, size_t size);
+  virtual size_t Write(const void *ptr, size_t size) override;
   // destructor
   virtual ~WriteStream() {
     this->Close();
@@ -863,13 +864,14 @@ class WriteStream : public Stream {
   void Finish(void);
 };
 
-void WriteStream::Write(const void *ptr, size_t size) {
+size_t WriteStream::Write(const void *ptr, size_t size) {
   size_t rlen = buffer_.length();
   buffer_.resize(rlen + size);
   std::memcpy(BeginPtr(buffer_) + rlen, ptr, size);
   if (buffer_.length() >= max_buffer_size_) {
     this->Upload();
   }
+  return size;
 }
 
 void WriteStream::Run(const std::string &method,
diff --git a/src/io/single_file_split.h b/src/io/single_file_split.h
@@ -70,8 +70,9 @@ class SingleFileSplit : public InputSplit {
     CHECK(part_index == 0 && num_parts == 1);
     this->BeforeFirst();
   }
-  virtual void Write(const void * /*ptr*/, size_t /*size*/) {
+  virtual size_t Write(const void * /*ptr*/, size_t /*size*/) {
     LOG(FATAL) << "InputSplit do not support write";
+    return 0;
   }
   virtual bool NextRecord(Blob *out_rec) {
     if (chunk_begin_ == chunk_end_) {

Original file line number	Diff line number	Diff line change
`@@ -30,23 +30,24 @@ struct MemoryFixedSizeStream : public SeekStream {`
`30`	`30`	`buffer_size_(buffer_size) {`
`31`	`31`	`curr_ptr_ = 0;`
`32`	`32`	`}`
`33`		`- virtual size_t Read(void *ptr, size_t size) {`
	`33`	`+ virtual size_t Read(void *ptr, size_t size) override {`
`34`	`34`	`CHECK(curr_ptr_ + size <= buffer_size_);`
`35`	`35`	`size_t nread = std::min(buffer_size_ - curr_ptr_, size);`
`36`	`36`	`if (nread != 0) std::memcpy(ptr, p_buffer_ + curr_ptr_, nread);`
`37`	`37`	`curr_ptr_ += nread;`
`38`	`38`	`return nread;`
`39`	`39`	`}`
`40`		`- virtual void Write(const void *ptr, size_t size) {`
`41`		`- if (size == 0) return;`
	`40`	`+ virtual size_t Write(const void *ptr, size_t size) override {`
	`41`	`+ if (size == 0) return 0;`
`42`	`42`	`CHECK(curr_ptr_ + size <= buffer_size_);`
`43`	`43`	`std::memcpy(p_buffer_ + curr_ptr_, ptr, size);`
`44`	`44`	`curr_ptr_ += size;`
	`45`	`+ return size;`
`45`	`46`	`}`
`46`		`- virtual void Seek(size_t pos) {`
	`47`	`+ virtual void Seek(size_t pos) override {`
`47`	`48`	`curr_ptr_ = static_cast<size_t>(pos);`
`48`	`49`	`}`
`49`		`- virtual size_t Tell(void) {`
	`50`	`+ virtual size_t Tell(void) override {`
`50`	`51`	`return curr_ptr_;`
`51`	`52`	`}`
`52`	`53`
`@@ -73,25 +74,26 @@ struct MemoryStringStream : public dmlc::SeekStream {`
`73`	`74`	`: p_buffer_(p_buffer) {`
`74`	`75`	`curr_ptr_ = 0;`
`75`	`76`	`}`
`76`		`- virtual size_t Read(void *ptr, size_t size) {`
	`77`	`+ virtual size_t Read(void *ptr, size_t size) override {`
`77`	`78`	`CHECK(curr_ptr_ <= p_buffer_->length());`
`78`	`79`	`size_t nread = std::min(p_buffer_->length() - curr_ptr_, size);`
`79`	`80`	`if (nread != 0) std::memcpy(ptr, &(*p_buffer_)[0] + curr_ptr_, nread);`
`80`	`81`	`curr_ptr_ += nread;`
`81`	`82`	`return nread;`
`82`	`83`	`}`
`83`		`- virtual void Write(const void *ptr, size_t size) {`
`84`		`- if (size == 0) return;`
	`84`	`+ virtual size_t Write(const void *ptr, size_t size) override {`
	`85`	`+ if (size == 0) return 0;`
`85`	`86`	`if (curr_ptr_ + size > p_buffer_->length()) {`
`86`	`87`	`p_buffer_->resize(curr_ptr_+size);`
`87`	`88`	`}`
`88`	`89`	`std::memcpy(&(*p_buffer_)[0] + curr_ptr_, ptr, size);`
`89`	`90`	`curr_ptr_ += size;`
	`91`	`+ return size;`
`90`	`92`	`}`
`91`		`- virtual void Seek(size_t pos) {`
	`93`	`+ virtual void Seek(size_t pos) override {`
`92`	`94`	`curr_ptr_ = static_cast<size_t>(pos);`
`93`	`95`	`}`
`94`		`- virtual size_t Tell(void) {`
	`96`	`+ virtual size_t Tell(void) override {`
`95`	`97`	`return curr_ptr_;`
`96`	`98`	`}`
`97`	`99`
Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,7 @@ class HDFSStream : public SeekStream {`
`28`	`28`	`}`
`29`	`29`	`}`
`30`	`30`
`31`		`- virtual size_t Read(void *ptr, size_t size) {`
	`31`	`+ virtual size_t Read(void *ptr, size_t size) override {`
`32`	`32`	`char buf = static_cast<char>(ptr);`
`33`	`33`	`size_t nleft = size;`
`34`	`34`	`size_t nmax = static_cast<size_t>(std::numeric_limits<tSize>::max());`
`@@ -48,7 +48,7 @@ class HDFSStream : public SeekStream {`
`48`	`48`	`return size - nleft;`
`49`	`49`	`}`
`50`	`50`
`51`		`- virtual void Write(const void *ptr, size_t size) {`
	`51`	`+ virtual size_t Write(const void *ptr, size_t size) override {`
`52`	`52`	`const char buf = reinterpret_cast<const char>(ptr);`
`53`	`53`	`size_t nleft = size;`
`54`	`54`	`// When using builtin-java classes to write, the maximum write size`
`@@ -70,14 +70,15 @@ class HDFSStream : public SeekStream {`
`70`	`70`	`LOG(FATAL) << "HDFSStream.hdfsWrite Error:" << strerror(errsv);`
`71`	`71`	`}`
`72`	`72`	`}`
	`73`	`+ return size - nleft;`
`73`	`74`	`}`
`74`		`- virtual void Seek(size_t pos) {`
	`75`	`+ virtual void Seek(size_t pos) override {`
`75`	`76`	`if (hdfsSeek(fs_, fp_, pos) != 0) {`
`76`	`77`	`int errsv = errno;`
`77`	`78`	`LOG(FATAL) << "HDFSStream.hdfsSeek Error:" << strerror(errsv);`
`78`	`79`	`}`
`79`	`80`	`}`
`80`		`- virtual size_t Tell(void) {`
	`81`	`+ virtual size_t Tell(void) override {`
`81`	`82`	`tOffset offset = hdfsTell(fs_, fp_);`
`82`	`83`	`if (offset == -1) {`
`83`	`84`	`int errsv = errno;`
Original file line number	Diff line number	Diff line change
`@@ -424,23 +424,24 @@ class CURLReadStreamBase : public SeekStream {`
`424`	`424`	`virtual ~CURLReadStreamBase() {`
`425`	`425`	`this->Cleanup();`
`426`	`426`	`}`
`427`		`- virtual size_t Tell(void) {`
	`427`	`+ virtual size_t Tell(void) override {`
`428`	`428`	`return curr_bytes_;`
`429`	`429`	`}`
`430`	`430`	`virtual bool AtEnd(void) const {`
`431`	`431`	`return at_end_;`
`432`	`432`	`}`
`433`		`- virtual void Write(const void *ptr, size_t size) {`
	`433`	`+ virtual size_t Write(const void *ptr, size_t size) override {`
`434`	`434`	`LOG(FATAL) << "CURL.ReadStream cannot be used for write";`
	`435`	`+ return 0;`
`435`	`436`	`}`
`436`	`437`	`// lazy seek function`
`437`		`- virtual void Seek(size_t pos) {`
	`438`	`+ virtual void Seek(size_t pos) override {`
`438`	`439`	`if (curr_bytes_ != pos) {`
`439`	`440`	`this->Cleanup();`
`440`	`441`	`curr_bytes_ = pos;`
`441`	`442`	`}`
`442`	`443`	`}`
`443`		`- virtual size_t Read(void *ptr, size_t size);`
	`444`	`+ virtual size_t Read(void *ptr, size_t size) override ;`
`444`	`445`
`445`	`446`	`protected:`
`446`	`447`	`CURLReadStreamBase()`
`@@ -790,11 +791,11 @@ class WriteStream : public Stream {`
`790`	`791`	`ecurl_ = curl_easy_init();`
`791`	`792`	`this->Init();`
`792`	`793`	`}`
`793`		`- virtual size_t Read(void *ptr, size_t size) {`
	`794`	`+ virtual size_t Read(void *ptr, size_t size) override {`
`794`	`795`	`LOG(FATAL) << "S3.WriteStream cannot be used for read";`
`795`	`796`	`return 0;`
`796`	`797`	`}`
`797`		`- virtual void Write(const void *ptr, size_t size);`
	`798`	`+ virtual size_t Write(const void *ptr, size_t size) override;`
`798`	`799`	`// destructor`
`799`	`800`	`virtual ~WriteStream() {`
`800`	`801`	`this->Close();`
`@@ -863,13 +864,14 @@ class WriteStream : public Stream {`
`863`	`864`	`void Finish(void);`
`864`	`865`	`};`
`865`	`866`
`866`		`-void WriteStream::Write(const void *ptr, size_t size) {`
	`867`	`+size_t WriteStream::Write(const void *ptr, size_t size) {`
`867`	`868`	`size_t rlen = buffer_.length();`
`868`	`869`	`buffer_.resize(rlen + size);`
`869`	`870`	`std::memcpy(BeginPtr(buffer_) + rlen, ptr, size);`
`870`	`871`	`if (buffer_.length() >= max_buffer_size_) {`
`871`	`872`	`this->Upload();`
`872`	`873`	`}`
	`874`	`+ return size;`
`873`	`875`	`}`
`874`	`876`
`875`	`877`	`void WriteStream::Run(const std::string &method,`
Original file line number	Diff line number	Diff line change
`@@ -70,8 +70,9 @@ class SingleFileSplit : public InputSplit {`
`70`	`70`	`CHECK(part_index == 0 && num_parts == 1);`
`71`	`71`	`this->BeforeFirst();`
`72`	`72`	`}`
`73`		`- virtual void Write(const void * /ptr/, size_t /size/) {`
	`73`	`+ virtual size_t Write(const void * /ptr/, size_t /size/) {`
`74`	`74`	`LOG(FATAL) << "InputSplit do not support write";`
	`75`	`+ return 0;`
`75`	`76`	`}`
`76`	`77`	`virtual bool NextRecord(Blob *out_rec) {`
`77`	`78`	`if (chunk_begin_ == chunk_end_) {`