root/trunk/whisperlib/common/io/logio/test/recordio_test.cc

// Copyright (c) 2009, Whispersoft s.r.l.
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions are
// met:
//
// * Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
// * Redistributions in binary form must reproduce the above
// copyright notice, this list of conditions and the following disclaimer
// in the documentation and/or other materials provided with the
// distribution.
// * Neither the name of Whispersoft s.r.l. nor the names of its
// contributors may be used to endorse or promote products derived from
// this software without specific prior written permission.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
//
// Author: Catalin Popescu

// Simple tests for recordio

#include "common/base/types.h"
#include "common/base/log.h"
#include "common/base/timer.h"
#include "common/base/system.h"

#include "common/io/buffer/memory_stream.h"
#include "common/io/logio/recordio.h"

//////////////////////////////////////////////////////////////////////

DEFINE_int32(num_records,
             1000,
             "Number of records for the test");

DEFINE_int32(rand_seed,
             3274,
             "Seed the random number generator w/ this");

DEFINE_int32(block_size,
             16384,
             "Create blocks of this size");

DEFINE_int32(max_record_size,
             65536,
             "Generate records of at most this size");

DEFINE_bool(deflate,
            false,
            "Zip records for writing");

//////////////////////////////////////////////////////////////////////

static unsigned int rand_seed;

bool GenerateRecord(io::RecordWriter* rw,
                    io::MemoryStream* rec,
                    io::MemoryStream* out) {
  const int32 rec_size = rand_r(&rand_seed) %
                         (FLAGS_max_record_size / sizeof(rec_size));
  int32* buf = new int32[rec_size];
  for ( int32 i = 0; i < rec_size; ++i ) {
    buf[i] = rand_r(&rand_seed);
  }
  rec->Write(buf, rec_size * sizeof(*buf));
  delete [] buf;
  rec->MarkerSet();
  bool ret = rw->AppendRecord(rec, out);
  rec->MarkerRestore();
  return ret;
}



int main(int argc, char* argv[]) {
  common::Init(argc, argv);

  rand_seed = FLAGS_rand_seed;
  srand(rand_seed);
  io::RecordWriter rw(FLAGS_block_size, FLAGS_deflate);
  vector<io::MemoryStream*> recs;
  io::MemoryStream out;

  for ( int32 i = 0; i < FLAGS_num_records; ++i ) {
    recs.push_back(new io::MemoryStream());
    if ( GenerateRecord(&rw, recs.back(), &out) ) {
      LOG_INFO << " Another block generated @" << out.Size()
               << " rec: " << i;
    }
    if ( recs.back()->IsEmpty() ) {
      LOG_INFO << "================ > EMPTY record at: " << i;
    }
  }
  rw.FinalizeContent(&out);
  LOG_INFO << " Final records size: " << out.Size();

  io::RecordReader rd(FLAGS_block_size);
  io::MemoryStream in;
  io::MemoryStream crt;
  int32 rec_id = 0;
  while ( !out.IsEmpty() ) {
    int32 to_pull = min(rand_r(&rand_seed) % 2048, out.Size());
    in.AppendStream(&out, to_pull);
    io::RecordReader::ReadResult err;
    do {
      err = rd.ReadRecord(&in, &crt);
      if ( err == io::RecordReader::READ_OK ) {
        CHECK_EQ(crt.Size(), recs[rec_id]->Size()) << " @ record: " << rec_id;
        string s1, s2;
        crt.ReadString(&s1);
        recs[rec_id]->ReadString(&s2);
        CHECK_EQ(s1, s2) << " @ record: " << rec_id;
        delete recs[rec_id];
        ++rec_id;
      }
      CHECK(err <= io::RecordReader::READ_NO_DATA) << err << " @ " << rec_id;
    } while ( err == io::RecordReader::READ_OK );
  }
  CHECK_EQ(rec_id, recs.size());

  printf("PASS\n");
}