folly/io/RecordIO.cpp

   1 /*
   2  * Copyright 2013 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #include "folly/io/RecordIO.h"
  18
  19 #include <sys/types.h>
  20 #include <unistd.h>
  21
  22 #include "folly/Exception.h"
  23 #include "folly/FileUtil.h"
  24 #include "folly/Memory.h"
  25 #include "folly/ScopeGuard.h"
  26 #include "folly/String.h"
  27
  28 namespace folly {
  29
  30 using namespace recordio_helpers;
  31
  32 RecordIOWriter::RecordIOWriter(File file, uint32_t fileId)
  33   : file_(std::move(file)),
  34     fileId_(fileId),
  35     writeLock_(file_, std::defer_lock),
  36     filePos_(0) {
  37   if (!writeLock_.try_lock()) {
  38     throw std::runtime_error("RecordIOWriter: file locked by another process");
  39   }
  40
  41   struct stat st;
  42   checkUnixError(fstat(file_.fd(), &st), "fstat() failed");
  43
  44   filePos_ = st.st_size;
  45 }
  46
  47 void RecordIOWriter::write(std::unique_ptr<IOBuf> buf) {
  48   size_t totalLength = prependHeader(buf, fileId_);
  49   if (totalLength == 0) {
  50     return;  // nothing to do
  51   }
  52
  53   DCHECK_EQ(buf->computeChainDataLength(), totalLength);
  54   auto iov = buf->getIov();
  55
  56   // We're going to write.  Reserve space for ourselves.
  57   off_t pos = filePos_.fetch_add(totalLength);
  58   ssize_t bytes = pwritevFull(file_.fd(), iov.data(), iov.size(), pos);
  59   checkUnixError(bytes, "pwrite() failed");
  60   DCHECK_EQ(bytes, totalLength);
  61 }
  62
  63 RecordIOReader::RecordIOReader(File file, uint32_t fileId)
  64   : map_(std::move(file)),
  65     fileId_(fileId) {
  66 }
  67
  68 RecordIOReader::Iterator::Iterator(ByteRange range, uint32_t fileId, off_t pos)
  69   : range_(range),
  70     fileId_(fileId),
  71     recordAndPos_(ByteRange(), 0) {
  72   if (pos >= range_.size()) {
  73     recordAndPos_.second = off_t(-1);
  74     range_.clear();
  75   } else {
  76     recordAndPos_.second = pos;
  77     range_.advance(pos);
  78     advanceToValid();
  79   }
  80 }
  81
  82 void RecordIOReader::Iterator::advanceToValid() {
  83   ByteRange record = findRecord(range_, fileId_).record;
  84   if (record.empty()) {
  85     recordAndPos_ = std::make_pair(ByteRange(), off_t(-1));
  86     range_.clear();  // at end
  87   } else {
  88     size_t skipped = record.begin() - range_.begin();
  89     DCHECK_GE(skipped, headerSize());
  90     skipped -= headerSize();
  91     range_.advance(skipped);
  92     recordAndPos_.first = record;
  93     recordAndPos_.second += skipped;
  94   }
  95 }
  96
  97 namespace recordio_helpers {
  98
  99 using namespace detail;
 100
 101 namespace {
 102
 103 constexpr uint32_t kHashSeed = 0xdeadbeef;  // for mcurtiss
 104
 105 uint32_t headerHash(const Header& header) {
 106   return hash::SpookyHashV2::Hash32(&header, offsetof(Header, headerHash),
 107                                     kHashSeed);
 108 }
 109
 110 std::pair<size_t, uint64_t> dataLengthAndHash(const IOBuf* buf) {
 111   size_t len = 0;
 112   hash::SpookyHashV2 hasher;
 113   hasher.Init(kHashSeed, kHashSeed);
 114   for (auto br : *buf) {
 115     len += br.size();
 116     hasher.Update(br.data(), br.size());
 117   }
 118   uint64_t hash1;
 119   uint64_t hash2;
 120   hasher.Final(&hash1, &hash2);
 121   if (len + headerSize() >= std::numeric_limits<uint32_t>::max()) {
 122     throw std::invalid_argument("Record length must fit in 32 bits");
 123   }
 124   return std::make_pair(len, hash1);
 125 }
 126
 127 uint64_t dataHash(ByteRange range) {
 128   return hash::SpookyHashV2::Hash64(range.data(), range.size(), kHashSeed);
 129 }
 130
 131 }  // namespace
 132
 133 size_t prependHeader(std::unique_ptr<IOBuf>& buf, uint32_t fileId) {
 134   if (fileId == 0) {
 135     throw std::invalid_argument("invalid file id");
 136   }
 137   auto lengthAndHash = dataLengthAndHash(buf.get());
 138   if (lengthAndHash.first == 0) {
 139     return 0;  // empty, nothing to do, no zero-length records
 140   }
 141
 142   // Prepend to the first buffer in the chain if we have room, otherwise
 143   // prepend a new buffer.
 144   if (buf->headroom() >= headerSize()) {
 145     buf->unshareOne();
 146     buf->prepend(headerSize());
 147   } else {
 148     auto b = IOBuf::create(headerSize());
 149     b->append(headerSize());
 150     b->appendChain(std::move(buf));
 151     buf = std::move(b);
 152   }
 153   detail::Header* header =
 154     reinterpret_cast<detail::Header*>(buf->writableData());
 155   memset(header, 0, sizeof(Header));
 156   header->magic = detail::Header::kMagic;
 157   header->fileId = fileId;
 158   header->dataLength = lengthAndHash.first;
 159   header->dataHash = lengthAndHash.second;
 160   header->headerHash = headerHash(*header);
 161
 162   return lengthAndHash.first + headerSize();
 163 }
 164
 165 RecordInfo validateRecord(ByteRange range, uint32_t fileId) {
 166   if (range.size() <= headerSize()) {  // records may not be empty
 167     return {0};
 168   }
 169   const Header* header = reinterpret_cast<const Header*>(range.begin());
 170   range.advance(sizeof(Header));
 171   if (header->magic != Header::kMagic ||
 172       header->version != 0 ||
 173       header->hashFunction != 0 ||
 174       header->flags != 0 ||
 175       (fileId != 0 && header->fileId != fileId) ||
 176       header->dataLength > range.size()) {
 177     return {0};
 178   }
 179   if (headerHash(*header) != header->headerHash) {
 180     return {0};
 181   }
 182   range.reset(range.begin(), header->dataLength);
 183   if (dataHash(range) != header->dataHash) {
 184     return {0};
 185   }
 186   return {header->fileId, range};
 187 }
 188
 189 RecordInfo findRecord(ByteRange searchRange,
 190                       ByteRange wholeRange,
 191                       uint32_t fileId) {
 192   static const uint32_t magic = Header::kMagic;
 193   static const ByteRange magicRange(reinterpret_cast<const uint8_t*>(&magic),
 194                                     sizeof(magic));
 195   static constexpr size_t headerTail = sizeof(Header) - sizeof(magic);
 196
 197   DCHECK_GE(searchRange.begin(), wholeRange.begin());
 198   DCHECK_LE(searchRange.end(), wholeRange.end());
 199
 200   const uint8_t* start = searchRange.begin();
 201   const uint8_t* end = std::min(searchRange.end(),
 202                                 wholeRange.end() - sizeof(Header));
 203   // end-1: the last place where a Header could start
 204   while (start < end) {
 205     auto p = ByteRange(start, end + sizeof(magic)).find(magicRange);
 206     if (p == ByteRange::npos) {
 207       break;
 208     }
 209
 210     start += p;
 211     auto r = validateRecord(ByteRange(start, wholeRange.end()), fileId);
 212     if (!r.record.empty()) {
 213       return r;
 214     }
 215
 216     // No repeated prefix in magic, so we can do better than start++
 217     start += sizeof(magic);
 218   }
 219
 220   return {0};
 221 }
 222
 223 }  // namespace
 224
 225 }  // namespaces