example/merge.cpp

   1 #include <stdlib.h>
   2 #include <string.h>
   3
   4 #include <fstream>
   5 #include <iostream>
   6 #include <string>
   7 #include <stdexcept>
   8 #include <vector>
   9
  10 #include <boost/shared_ptr.hpp>
  11
  12 #include <eblob/blob.h>
  13 #include "common.hpp"
  14
  15 static void copy_data(std::ifstream &src, std::ofstream &dst, size_t size)
  16 {
  17         size_t sz = 1024 * 1024;
  18
  19         char *buf = new char[sz];
  20
  21         try {
  22                 while (size != 0) {
  23                         size_t tmp = sz;
  24                         if (tmp > size)
  25                                 tmp = size;
  26
  27                         src.read(buf, tmp);
  28                         dst.write(buf, tmp);
  29
  30                         size -= tmp;
  31                 }
  32         } catch (...) {
  33                 delete [] buf;
  34                 throw;
  35         }
  36
  37         delete [] buf;
  38 }
  39
  40 static void em_usage(char *p)
  41 {
  42         std::cerr << "Usage: " << p << "<options>" << std::endl <<
  43                 "  This utility will defragment and merge (multiple) blobs into larger one\n"
  44                 "  -i path             - input blob path (can be specified multiple times)\n"
  45                 "  -o path             - output blob path\n"
  46                 "  -p                  - print all copied IDs\n"
  47                 "  -h                  - this help\n"
  48                 "" << std::endl;
  49         exit(-1);
  50 }
  51
  52 struct em_blob {
  53         int                                     completed;
  54         std::ifstream                           index, data;
  55         std::string                             path_;
  56
  57         em_blob(const char *path) : completed(0), path_(path) {
  58                 try {
  59                         data.open(path, std::ios_base::in | std::ios_base::binary);
  60                         std::string index_path(path);
  61                         index_path += ".index";
  62                         index.open(index_path.c_str(), std::ios_base::in | std::ios_base::binary);
  63                 } catch (...) {
  64                         data.close();
  65                         index.close();
  66
  67                         throw;
  68                 }
  69         }
  70
  71         em_blob(const struct em_blob &e) {
  72                 em_blob(e.path_.c_str());
  73         }
  74
  75         ~em_blob() {
  76                 data.close();
  77                 index.close();
  78         }
  79 };
  80
  81 typedef boost::shared_ptr<em_blob> em_blob_ptr;
  82
  83 struct em_ctl {
  84         struct eblob_disk_control               dc;
  85         em_blob_ptr                             blob;
  86
  87         em_ctl(em_blob_ptr b) : blob(b) {
  88                 memset(&dc, 0, sizeof(struct eblob_disk_control));
  89         }
  90 };
  91
  92 struct em_compare {
  93                 bool operator () (const em_ctl &s1, const em_ctl &s2) const {
  94                         return memcmp(s1.dc.key.id, s2.dc.key.id, EBLOB_ID_SIZE);
  95                 }
  96 };
  97
  98 int main(int argc, char *argv[])
  99 {
 100         int ch;
 101         int total_input = 0;
 102         int print_all = 0;
 103         struct eblob_disk_control ddc;
 104         long long total = 0, removed = 0, written = 0, broken = 0;
 105         long long position = 0;
 106
 107         std::vector<em_blob_ptr> blobs;
 108         std::string output;
 109
 110         while ((ch = getopt(argc, argv, "i:o:ph")) != -1) {
 111                 switch (ch) {
 112                         case 'i':
 113                                 try {
 114                                         em_blob_ptr b(new em_blob(optarg));
 115
 116                                         blobs.push_back(b);
 117                                         total_input++;
 118                                 } catch (const std::exception &e) {
 119                                         std::cerr << "could not open data or index file for blob " << optarg << ": " << e.what() << std::endl;
 120                                 }
 121                                 break;
 122                         case 'o':
 123                                 output.assign(optarg);
 124                                 break;
 125                         case 'p':
 126                                 print_all = 1;
 127                                 break;
 128                         case 'h':
 129                         default:
 130                                 em_usage(argv[0]);
 131                                 /* not reached */
 132                 }
 133         }
 134
 135         if (!blobs.size() || !output.size()) {
 136                 std::cerr << "You must specify input and output parameters" << std::endl;
 137                 em_usage(argv[0]);
 138         }
 139
 140         try {
 141                 std::string data_path = output;
 142                 std::string index_path = output + ".index";
 143
 144                 std::ofstream index_out(index_path.c_str(), std::ios_base::out | std::ios_base::binary | std::ios::trunc);
 145                 std::ofstream data_out(data_path.c_str(), std::ios_base::out | std::ios_base::binary | std::ios::trunc);
 146
 147                 while (blobs.size() != 0) {
 148                         std::vector<struct em_ctl> ctl;
 149
 150                         for (std::vector<em_blob_ptr>::iterator b = blobs.begin(); b < blobs.end(); ++b) {
 151                                 struct em_ctl c(*b);
 152
 153                                 struct em_blob *blob = b->get();
 154
 155                                 if (blob->completed)
 156                                         continue;
 157
 158                                 do {
 159                                         blob->index.read((char *)&c.dc, sizeof(struct eblob_disk_control));
 160
 161                                         if (blob->index.gcount() != sizeof(struct eblob_disk_control)) {
 162                                                 blob->completed = 1;
 163
 164                                                 std::cout << "Completed input stream " << blob->path_ <<
 165                                                         ": total: " << total_input <<
 166                                                         ", rest: " << blobs.size() << std::endl;
 167                                                 break;
 168
 169                                         }
 170                                 } while (c.dc.disk_size == 0);
 171
 172                                 if (blob->completed)
 173                                         continue;
 174
 175                                 blob->index.seekg(-sizeof(struct eblob_disk_control), std::ios_base::cur);
 176
 177                                 eblob_convert_disk_control(&c.dc);
 178                                 ctl.push_back(c);
 179                         }
 180
 181                         if (!ctl.size()) {
 182                                 std::cout << "Completed all blobs" << std::endl;
 183                                 break;
 184                         }
 185
 186
 187                         std::sort(ctl.begin(), ctl.end(), em_compare());
 188
 189                         total++;
 190
 191                         struct em_ctl c = ctl[0];
 192                         c.blob->index.seekg(sizeof(struct eblob_disk_control), std::ios_base::cur);
 193
 194                         if (print_all) {
 195                                 std::cout << c.blob->path_ << ": " << eblob_dump_control(&c.dc, position, 1, 0) << std::endl;
 196                         }
 197
 198                         if (c.dc.flags & BLOB_DISK_CTL_REMOVE) {
 199                                 removed++;
 200                                 continue;
 201                         }
 202
 203                         c.blob->data.seekg(c.dc.position, std::ios::beg);
 204                         c.blob->data.read((char *)&ddc, sizeof(struct eblob_disk_control));
 205                         eblob_convert_disk_control(&ddc);
 206
 207                         if (c.blob->data.gcount() != sizeof(struct eblob_disk_control))
 208                                 throw std::runtime_error("Data read failed");
 209
 210                         if (ddc.flags & BLOB_DISK_CTL_REMOVE) {
 211                                 removed++;
 212                                 continue;
 213                         }
 214
 215                         size_t size = ddc.disk_size;
 216
 217                         ddc.position = position;
 218
 219                         if (print_all) {
 220                                 std::cout << "out: " << eblob_dump_control(&ddc, position, 1, 0) << std::endl;
 221                         }
 222
 223                         if (size > sizeof(struct eblob_disk_control)) {
 224                                 eblob_convert_disk_control(&ddc);
 225
 226                                 data_out.write((char *)&ddc, sizeof(struct eblob_disk_control));
 227                                 copy_data(c.blob->data, data_out, size - sizeof(struct eblob_disk_control));
 228
 229                                 index_out.write((char *)&ddc, sizeof(struct eblob_disk_control));
 230
 231                                 position += size;
 232                                 written++;
 233                         } else {
 234                                 broken++;
 235                         }
 236                 }
 237         } catch (const std::exception &e) {
 238                 std::cerr << e.what() << std::endl;
 239         }
 240
 241         std::cout << "Total records: " << total << std::endl;
 242         std::cout << "Written records: " << written << std::endl;
 243         std::cout << "Removed records: " << removed << std::endl;
 244         std::cout << "Broken records: " << broken << std::endl;
 245
 246         return 0;
 247 }
 248