package com.test.mongodb;

import java.io.File;  
import java.io.FileOutputStream;  
import java.io.IOException;  
import java.util.UUID;
import java.util.concurrent.atomic.AtomicLong;  
  
import org.apache.jmeter.config.Arguments;  
import org.apache.jmeter.protocol.java.sampler.AbstractJavaSamplerClient;  
import org.apache.jmeter.protocol.java.sampler.JavaSamplerContext;  
import org.apache.jmeter.samplers.SampleResult;  
  
import com.mongodb.BasicDBObject;  
import com.mongodb.DB;  
import com.mongodb.DBCollection;  
import com.mongodb.Mongo;  
import com.mongodb.MongoOptions;  
import com.mongodb.ServerAddress;  
import com.mongodb.WriteResult;  
  
/** 
 *  
 * @author  2014-10-14
 * @remark  一次性插入多条数据测试MongoDB的性能    
 */  
  
public class TestMongodb extends AbstractJavaSamplerClient {  
  
    private static Mongo            m;  
    private static DB               db;  
    private static AtomicLong       read_key = new AtomicLong(0);  
    private static FileOutputStream fos;  
    static {  
        try {  
            fos = new FileOutputStream(new File("jmeter_error.log"));  
  
            MongoOptions options = new MongoOptions();  
            options.autoConnectRetry = true;  
            options.connectionsPerHost = 1000;  
            options.maxWaitTime = 5000;  
            options.socketTimeout = 0;  
            options.connectTimeout = 15000;  
            options.threadsAllowedToBlockForConnectionMultiplier = 5000;  
            m = new Mongo(new ServerAddress("127.0.0.1", 27017), options);  
            db = m.getDB("iflashbuy-log");  
        } catch (Exception e) {  
            throw new RuntimeException(e);  
        }  
    }  
  
    public Arguments getDefaultParameters() {
        Arguments params = new Arguments();
        params.addArgument("rw", "w");
        return params;
    }
    
    public SampleResult runTest(JavaSamplerContext context) {  
        // System.out.println(read_key.getAndIncrement());  
        SampleResult results = new SampleResult();  
        //System.out.println(context.getParameter("rw"));  
        results.sampleStart();  
        DBCollection coll = db.getCollection("area");  
        if ("w".equalsIgnoreCase(context.getParameter("rw"))) {  
            long key = 1;  
            BasicDBObject doc = new BasicDBObject();                
            doc.put("_id", UUID.randomUUID());  
            doc.put("_class", "com.iflashbuy.bi.source.pojo.AreaDataEntity");  
            doc.put("province", "广东");  
            doc.put("city", "广州");  
            doc.put("pv", 35522924);  
            doc.put("uv", 52556);  
            doc.put("orderCount", 963);  
            doc.put("orderTotal", 1548563);  
            doc.put("entpriseCode", "00540001");  
            db.requestStart();  
            WriteResult rs = coll.insert(doc);  
            try {  
                if (rs.getError() != null) fos.write(rs.getError().getBytes());  
            } catch (IOException e) {  
                // TODO Auto-generated catch block  
                e.printStackTrace();  
            }  
            //System.out.println("w error->" + rs.getError());  
            db.requestDone();  
        }  

        results.setSuccessful(true);  
        results.sampleEnd();  
        return results;  
        // return null;  
    }  
  
    /*
    public static void main(String args[]) throws Exception {  
        Arguments a = new Arguments();  
        a.addArgument("rw", "w");  
        final JavaSamplerContext c = new JavaSamplerContext(a);  
        long start = System.currentTimeMillis();  
        for (int i = 0; i < 5; i++) {  
            new Thread() {  
                public void run() {  
                    TestMongodb t = new TestMongodb();  
                    for (int j = 0; j < 2; j++) {  
                        t.runTest(c);  
                    }  
                }  
            }.start();  
        }  
    }  
    */
}