1   /*
2    * Licensed to the Apache Software Foundation (ASF) under one or more
3    * contributor license agreements. See the NOTICE file distributed with this
4    * work for additional information regarding copyright ownership. The ASF
5    * licenses this file to you under the Apache License, Version 2.0 (the
6    * "License"); you may not use this file except in compliance with the License.
7    * You may obtain a copy of the License at
8    *
9    * http://www.apache.org/licenses/LICENSE-2.0
10   *
11   * Unless required by applicable law or agreed to in writing, software
12   * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
13   * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
14   * License for the specific language governing permissions and limitations
15   * under the License.
16   */
17  package org.apache.hadoop.hbase.util;
18  
19  import java.io.IOException;
20  import java.util.HashSet;
21  import java.util.PriorityQueue;
22  import java.util.Queue;
23  import java.util.Set;
24  import java.util.concurrent.ArrayBlockingQueue;
25  import java.util.concurrent.BlockingQueue;
26  import java.util.concurrent.ConcurrentSkipListSet;
27  import java.util.concurrent.TimeUnit;
28  import java.util.concurrent.atomic.AtomicLong;
29  
30  import org.apache.commons.logging.Log;
31  import org.apache.commons.logging.LogFactory;
32  import org.apache.hadoop.conf.Configuration;
33  import org.apache.hadoop.hbase.client.HTable;
34  import org.apache.hadoop.hbase.client.Put;
35  
36  /** Creates multiple threads that write key/values into the */
37  public class MultiThreadedWriter extends MultiThreadedAction {
38    private static final Log LOG = LogFactory.getLog(MultiThreadedWriter.class);
39  
40    private Set<HBaseWriterThread> writers = new HashSet<HBaseWriterThread>();
41  
42    private boolean isMultiPut = false;
43  
44    /**
45     * A temporary place to keep track of inserted keys. This is written to by
46     * all writers and is drained on a separate thread that populates
47     * {@link #insertedUpToKey}, the maximum key in the contiguous range of keys
48     * being inserted. This queue is supposed to stay small.
49     */
50    private BlockingQueue<Long> insertedKeys = new ArrayBlockingQueue<Long>(10000);
51  
52    /**
53     * This is the current key to be inserted by any thread. Each thread does an
54     * atomic get and increment operation and inserts the current value.
55     */
56    private AtomicLong nextKeyToInsert = new AtomicLong();
57  
58    /**
59     * The highest key in the contiguous range of keys .
60     */
61    private AtomicLong insertedUpToKey = new AtomicLong();
62  
63    /** The sorted set of keys NOT inserted by the writers */
64    private Set<Long> failedKeySet = new ConcurrentSkipListSet<Long>();
65  
66    /**
67     * The total size of the temporary inserted key set that have not yet lined
68     * up in a our contiguous sequence starting from startKey. Supposed to stay
69     * small.
70     */
71    private AtomicLong insertedKeyQueueSize = new AtomicLong();
72  
73    /** Enable this if used in conjunction with a concurrent reader. */
74    private boolean trackInsertedKeys;
75  
76    public MultiThreadedWriter(LoadTestDataGenerator dataGen, Configuration conf,
77        byte[] tableName) {
78      super(dataGen, conf, tableName, "W");
79    }
80  
81    /** Use multi-puts vs. separate puts for every column in a row */
82    public void setMultiPut(boolean isMultiPut) {
83      this.isMultiPut = isMultiPut;
84    }
85  
86    @Override
87    public void start(long startKey, long endKey, int numThreads)
88        throws IOException {
89      super.start(startKey, endKey, numThreads);
90  
91      if (verbose) {
92        LOG.debug("Inserting keys [" + startKey + ", " + endKey + ")");
93      }
94  
95      nextKeyToInsert.set(startKey);
96      insertedUpToKey.set(startKey - 1);
97  
98      for (int i = 0; i < numThreads; ++i) {
99        HBaseWriterThread writer = new HBaseWriterThread(i);
100       writers.add(writer);
101     }
102 
103     if (trackInsertedKeys) {
104       new Thread(new InsertedKeysTracker()).start();
105       numThreadsWorking.incrementAndGet();
106     }
107 
108     startThreads(writers);
109   }
110 
111   private class HBaseWriterThread extends Thread {
112     private final HTable table;
113 
114     public HBaseWriterThread(int writerId) throws IOException {
115       setName(getClass().getSimpleName() + "_" + writerId);
116       table = new HTable(conf, tableName);
117     }
118 
119     public void run() {
120       try {
121         long rowKeyBase;
122         byte[][] columnFamilies = dataGenerator.getColumnFamilies();
123         while ((rowKeyBase = nextKeyToInsert.getAndIncrement()) < endKey) {
124           byte[] rowKey = dataGenerator.getDeterministicUniqueKey(rowKeyBase);
125           Put put = new Put(rowKey);
126           numKeys.addAndGet(1);
127           int columnCount = 0;
128           for (byte[] cf : columnFamilies) {
129             byte[][] columns = dataGenerator.generateColumnsForCf(rowKey, cf);
130             for (byte[] column : columns) {
131               byte[] value = dataGenerator.generateValue(rowKey, cf, column);
132               put.add(cf, column, value);
133               ++columnCount;
134               if (!isMultiPut) {
135                 insert(table, put, rowKeyBase);
136                 numCols.addAndGet(1);
137                 put = new Put(rowKey);
138               }
139             }
140           }
141           if (isMultiPut) {
142             if (verbose) {
143               LOG.debug("Preparing put for key = [" + rowKey + "], " + columnCount + " columns");
144             }
145             insert(table, put, rowKeyBase);
146             numCols.addAndGet(columnCount);
147           }
148           if (trackInsertedKeys) {
149             insertedKeys.add(rowKeyBase);
150           }
151         }
152       } finally {
153         try {
154           table.close();
155         } catch (IOException e) {
156           LOG.error("Error closing table", e);
157         }
158         numThreadsWorking.decrementAndGet();
159       }
160     }
161   }
162 
163   public void insert(HTable table, Put put, long keyBase) {
164     try {
165       long start = System.currentTimeMillis();
166       table.put(put);
167       totalOpTimeMs.addAndGet(System.currentTimeMillis() - start);
168     } catch (IOException e) {
169       failedKeySet.add(keyBase);
170       LOG.error("Failed to insert: " + keyBase);
171       e.printStackTrace();
172     }
173   }
174 
175   /**
176    * A thread that keeps track of the highest key in the contiguous range of
177    * inserted keys.
178    */
179   private class InsertedKeysTracker implements Runnable {
180 
181     @Override
182     public void run() {
183       Thread.currentThread().setName(getClass().getSimpleName());
184       try {
185         long expectedKey = startKey;
186         Queue<Long> sortedKeys = new PriorityQueue<Long>();
187         while (expectedKey < endKey) {
188           // Block until a new element is available.
189           Long k;
190           try {
191             k = insertedKeys.poll(1, TimeUnit.SECONDS);
192           } catch (InterruptedException e) {
193             LOG.info("Inserted key tracker thread interrupted", e);
194             break;
195           }
196           if (k == null) {
197             continue;
198           }
199           if (k == expectedKey) {
200             // Skip the "sorted key" queue and consume this key.
201             insertedUpToKey.set(k);
202             ++expectedKey;
203           } else {
204             sortedKeys.add(k);
205           }
206 
207           // See if we have a sequence of contiguous keys lined up.
208           while (!sortedKeys.isEmpty()
209               && ((k = sortedKeys.peek()) == expectedKey)) {
210             sortedKeys.poll();
211             insertedUpToKey.set(k);
212             ++expectedKey;
213           }
214 
215           insertedKeyQueueSize.set(insertedKeys.size() + sortedKeys.size());
216         }
217       } catch (Exception ex) {
218         LOG.error("Error in inserted key tracker", ex);
219       } finally {
220         numThreadsWorking.decrementAndGet();
221       }
222     }
223 
224   }
225 
226   @Override
227   public void waitForFinish() {
228     super.waitForFinish();
229     System.out.println("Failed to write keys: " + failedKeySet.size());
230     for (Long key : failedKeySet) {
231        System.out.println("Failed to write key: " + key);
232     }
233   }
234 
235   public int getNumWriteFailures() {
236     return failedKeySet.size();
237   }
238 
239   /**
240    * The max key until which all keys have been inserted (successfully or not).
241    * @return the last key that we have inserted all keys up to (inclusive)
242    */
243   public long insertedUpToKey() {
244     return insertedUpToKey.get();
245   }
246 
247   public boolean failedToWriteKey(long k) {
248     return failedKeySet.contains(k);
249   }
250 
251   @Override
252   protected String progressInfo() {
253     StringBuilder sb = new StringBuilder();
254     appendToStatus(sb, "insertedUpTo", insertedUpToKey.get());
255     appendToStatus(sb, "insertedQSize", insertedKeyQueueSize.get());
256     return sb.toString();
257   }
258 
259   /**
260    * Used for a joint write/read workload. Enables tracking the last inserted
261    * key, which requires a blocking queue and a consumer thread.
262    * @param enable whether to enable tracking the last inserted key
263    */
264   public void setTrackInsertedKeys(boolean enable) {
265     trackInsertedKeys = enable;
266   }
267 
268 }