git.cworth.org Git - apitrace/blob - trace_file.cpp

   1 #include "trace_file.hpp"
   2
   3 #include <assert.h>
   4 #include <string.h>
   5
   6 #include <zlib.h>
   7 #include <snappy.h>
   8
   9 #include "os.hpp"
  10
  11 #include <iostream>
  12 #include <set>
  13
  14 using namespace Trace;
  15
  16 #define SNAPPY_BYTE1 'a'
  17 #define SNAPPY_BYTE2 't'
  18
  19 static void cleanupHandler(int sig);
  20
  21 class FileCleanup
  22 {
  23 public:
  24     FileCleanup()
  25         : m_quitOnFlush(false)
  26     {
  27         OS::CatchInterrupts(cleanupHandler);
  28     }
  29
  30     ~FileCleanup()
  31     {
  32     }
  33
  34     void addFile(Trace::File *file)
  35     {
  36         m_files.insert(file);
  37     }
  38     void removeFile(Trace::File *file)
  39     {
  40         m_files.erase(file);
  41         if (m_files.empty() && m_quitOnFlush) {
  42             OS::Abort();
  43         }
  44     }
  45
  46     void setQuitOnFlush(bool quit) {
  47         m_quitOnFlush = quit;
  48     }
  49     bool quitOnFlush() const
  50     {
  51         return m_quitOnFlush;
  52     }
  53
  54 private:
  55     std::set<Trace::File*> m_files;
  56     bool m_quitOnFlush;
  57 };
  58 static FileCleanup s_cleaner;
  59
  60 static void cleanupHandler(int sig)
  61 {
  62     s_cleaner.setQuitOnFlush(true);
  63 }
  64
  65 File::File(const std::string &filename,
  66            File::Mode mode)
  67     : m_filename(filename),
  68       m_mode(mode),
  69       m_isOpened(false)
  70 {
  71     if (!m_filename.empty()) {
  72         open(m_filename, m_mode);
  73     }
  74 }
  75
  76
  77 File::~File()
  78 {
  79     close();
  80 }
  81
  82 bool File::isOpened() const
  83 {
  84     return m_isOpened;
  85 }
  86
  87 File::Mode File::mode() const
  88 {
  89     return m_mode;
  90 }
  91
  92 std::string File::filename() const
  93 {
  94     return m_filename;
  95 }
  96
  97 bool File::open(const std::string &filename, File::Mode mode)
  98 {
  99     if (m_isOpened) {
 100         close();
 101     }
 102     m_isOpened = rawOpen(filename, mode);
 103     m_mode = mode;
 104
 105     if (m_isOpened) {
 106         s_cleaner.addFile(this);
 107     }
 108     return m_isOpened;
 109 }
 110
 111 bool File::write(const void *buffer, int length)
 112 {
 113     if (!m_isOpened || m_mode != File::Write) {
 114         return false;
 115     }
 116     return rawWrite(buffer, length);
 117 }
 118
 119 bool File::read(void *buffer, int length)
 120 {
 121     if (!m_isOpened || m_mode != File::Read) {
 122         return false;
 123     }
 124     return rawRead(buffer, length);
 125 }
 126
 127 void File::close()
 128 {
 129     if (m_isOpened) {
 130         rawClose();
 131         m_isOpened = false;
 132         s_cleaner.removeFile(this);
 133     }
 134 }
 135
 136 void File::flush()
 137 {
 138     rawFlush();
 139     if (s_cleaner.quitOnFlush()) {
 140         close();
 141     }
 142 }
 143
 144 int File::getc()
 145 {
 146     if (!m_isOpened || m_mode != File::Read) {
 147         return 0;
 148     }
 149     return rawGetc();
 150 }
 151
 152 bool File::isZLibCompressed(const std::string &filename)
 153 {
 154     std::fstream stream(filename.c_str(),
 155                         std::fstream::binary | std::fstream::in);
 156     if (!stream.is_open())
 157         return false;
 158
 159     unsigned char byte1, byte2;
 160     stream >> byte1;
 161     stream >> byte2;
 162     stream.close();
 163
 164     return (byte1 == 0x1f && byte2 == 0x8b);
 165 }
 166
 167
 168 bool File::isSnappyCompressed(const std::string &filename)
 169 {
 170     std::fstream stream(filename.c_str(),
 171                         std::fstream::binary | std::fstream::in);
 172     if (!stream.is_open())
 173         return false;
 174
 175     unsigned char byte1, byte2;
 176     stream >> byte1;
 177     stream >> byte2;
 178     stream.close();
 179
 180     return (byte1 == SNAPPY_BYTE1 && byte2 == SNAPPY_BYTE2);
 181 }
 182
 183
 184 ZLibFile::ZLibFile(const std::string &filename,
 185                    File::Mode mode)
 186     : File(filename, mode),
 187       m_gzFile(NULL)
 188 {
 189 }
 190
 191 ZLibFile::~ZLibFile()
 192 {
 193 }
 194
 195 bool ZLibFile::rawOpen(const std::string &filename, File::Mode mode)
 196 {
 197     m_gzFile = gzopen(filename.c_str(),
 198                       (mode == File::Write) ? "wb" : "rb");
 199     return m_gzFile != NULL;
 200 }
 201
 202 bool ZLibFile::rawWrite(const void *buffer, int length)
 203 {
 204     return gzwrite(m_gzFile, buffer, length) != -1;
 205 }
 206
 207 bool ZLibFile::rawRead(void *buffer, int length)
 208 {
 209     return gzread(m_gzFile, buffer, length) != -1;
 210 }
 211
 212 int ZLibFile::rawGetc()
 213 {
 214     return gzgetc(m_gzFile);
 215 }
 216
 217 void ZLibFile::rawClose()
 218 {
 219     if (m_gzFile) {
 220         gzclose(m_gzFile);
 221         m_gzFile = NULL;
 222     }
 223 }
 224
 225 void ZLibFile::rawFlush()
 226 {
 227     gzflush(m_gzFile, Z_SYNC_FLUSH);
 228 }
 229
 230 SnappyFile::SnappyFile(const std::string &filename,
 231                               File::Mode mode)
 232     : File(),
 233       m_cache(0),
 234       m_cachePtr(0),
 235       m_cacheSize(0)
 236 {
 237     m_compressedCache = new char[SNAPPY_CHUNK_SIZE];
 238 }
 239
 240 SnappyFile::~SnappyFile()
 241 {
 242     delete [] m_compressedCache;
 243 }
 244
 245 bool SnappyFile::rawOpen(const std::string &filename, File::Mode mode)
 246 {
 247     std::ios_base::openmode fmode = std::fstream::binary;
 248     if (mode == File::Write) {
 249         fmode |= (std::fstream::out | std::fstream::trunc);
 250         createCache(SNAPPY_CHUNK_SIZE);
 251     } else if (mode == File::Read) {
 252         fmode |= std::fstream::in;
 253     }
 254
 255     m_stream.open(filename.c_str(), fmode);
 256
 257     //read in the initial buffer if we're reading
 258     if (m_stream.is_open() && mode == File::Read) {
 259         // read the snappy file identifier
 260         unsigned char byte1, byte2;
 261         m_stream >> byte1;
 262         m_stream >> byte2;
 263         assert(byte1 == SNAPPY_BYTE1 && byte2 == SNAPPY_BYTE2);
 264
 265         flushCache();
 266     } else if (m_stream.is_open() && mode == File::Write) {
 267         // write the snappy file identifier
 268         m_stream << SNAPPY_BYTE1;
 269         m_stream << SNAPPY_BYTE2;
 270     }
 271     return m_stream.is_open();
 272 }
 273
 274 bool SnappyFile::rawWrite(const void *buffer, int length)
 275 {
 276     if (freeCacheSize() > length) {
 277         memcpy(m_cachePtr, buffer, length);
 278         m_cachePtr += length;
 279     } else if (freeCacheSize() == length) {
 280         memcpy(m_cachePtr, buffer, length);
 281         m_cachePtr += length;
 282         flushCache();
 283     } else {
 284         int sizeToWrite = length;
 285
 286         while (sizeToWrite >= freeCacheSize()) {
 287             int endSize = freeCacheSize();
 288             int offset = length - sizeToWrite;
 289             memcpy(m_cachePtr, (char*)buffer + offset, endSize);
 290             sizeToWrite -= endSize;
 291             m_cachePtr += endSize;
 292             flushCache();
 293         }
 294         if (sizeToWrite) {
 295             int offset = length - sizeToWrite;
 296             memcpy(m_cachePtr, (char*)buffer + offset, sizeToWrite);
 297             m_cachePtr += sizeToWrite;
 298         }
 299     }
 300
 301     return true;
 302 }
 303
 304 bool SnappyFile::rawRead(void *buffer, int length)
 305 {
 306     if (m_stream.eof()) {
 307         return false;
 308     }
 309     if (freeCacheSize() > length) {
 310         memcpy(buffer, m_cachePtr, length);
 311         m_cachePtr += length;
 312     } else if (freeCacheSize() == length) {
 313         memcpy(buffer, m_cachePtr, length);
 314         m_cachePtr += length;
 315         flushCache();
 316     } else {
 317         int sizeToRead = length;
 318         int offset = 0;
 319         while (sizeToRead) {
 320             int chunkSize = std::min(freeCacheSize(), sizeToRead);
 321             offset = length - sizeToRead;
 322             memcpy((char*)buffer + offset, m_cachePtr, chunkSize);
 323             m_cachePtr += chunkSize;
 324             sizeToRead -= chunkSize;
 325             if (sizeToRead > 0)
 326                 flushCache();
 327         }
 328     }
 329
 330     return true;
 331 }
 332
 333 int SnappyFile::rawGetc()
 334 {
 335     int c = 0;
 336     if (!rawRead(&c, 1))
 337         return -1;
 338     return c;
 339 }
 340
 341 void SnappyFile::rawClose()
 342 {
 343     flushCache();
 344     m_stream.close();
 345     delete [] m_cache;
 346     m_cache = NULL;
 347     m_cachePtr = NULL;
 348 }
 349
 350 void SnappyFile::rawFlush()
 351 {
 352     m_stream.flush();
 353 }
 354
 355 void SnappyFile::flushCache()
 356 {
 357     if (m_mode == File::Write) {
 358         size_t compressedLength;
 359
 360         ::snappy::RawCompress(m_cache, SNAPPY_CHUNK_SIZE - freeCacheSize(),
 361                               m_compressedCache, &compressedLength);
 362
 363         m_stream << compressedLength;
 364         m_stream.write(m_compressedCache, compressedLength);
 365         m_cachePtr = m_cache;
 366     } else if (m_mode == File::Read) {
 367         if (m_stream.eof())
 368             return;
 369         //assert(m_cachePtr == m_cache + m_cacheSize);
 370         size_t compressedLength;
 371         m_stream >> compressedLength;
 372         m_stream.read((char*)m_compressedCache, compressedLength);
 373         ::snappy::GetUncompressedLength(m_compressedCache, compressedLength,
 374                                         &m_cacheSize);
 375         if (m_cache)
 376             delete [] m_cache;
 377         createCache(m_cacheSize);
 378         ::snappy::RawUncompress(m_compressedCache, compressedLength,
 379                                 m_cache);
 380     }
 381 }
 382
 383 void SnappyFile::createCache(size_t size)
 384 {
 385     m_cache = new char[size];
 386     m_cachePtr = m_cache;
 387     m_cacheSize = size;
 388 }