single.go

Documentation: go.mongodb.org/mongo-driver/benchmark

     1  // Copyright (C) MongoDB, Inc. 2017-present.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License"); you may
     4  // not use this file except in compliance with the License. You may obtain
     5  // a copy of the License at http://www.apache.org/licenses/LICENSE-2.0
     6  
     7  package benchmark
     8  
     9  import (
    10  	"context"
    11  	"errors"
    12  
    13  	"go.mongodb.org/mongo-driver/bson"
    14  	"go.mongodb.org/mongo-driver/internal/handshake"
    15  	"go.mongodb.org/mongo-driver/internal/integtest"
    16  	"go.mongodb.org/mongo-driver/mongo"
    17  	"go.mongodb.org/mongo-driver/mongo/options"
    18  )
    19  
    20  const (
    21  	singleAndMultiDataDir = "single_and_multi_document"
    22  	tweetData             = "tweet.json"
    23  	smallData             = "small_doc.json"
    24  	largeData             = "large_doc.json"
    25  )
    26  
    27  func getClientDB(ctx context.Context) (*mongo.Database, error) {
    28  	cs, err := integtest.GetConnString()
    29  	if err != nil {
    30  		return nil, err
    31  	}
    32  	client, err := mongo.NewClient(options.Client().ApplyURI(cs.String()))
    33  	if err != nil {
    34  		return nil, err
    35  	}
    36  	if err = client.Connect(ctx); err != nil {
    37  		return nil, err
    38  	}
    39  
    40  	db := client.Database(integtest.GetDBName(cs))
    41  	return db, nil
    42  }
    43  
    44  func SingleRunCommand(ctx context.Context, tm TimerManager, iters int) error {
    45  	ctx, cancel := context.WithCancel(ctx)
    46  	defer cancel()
    47  
    48  	db, err := getClientDB(ctx)
    49  	if err != nil {
    50  		return err
    51  	}
    52  	defer func() { _ = db.Client().Disconnect(ctx) }()
    53  
    54  	cmd := bson.D{{handshake.LegacyHelloLowercase, true}}
    55  
    56  	tm.ResetTimer()
    57  	for i := 0; i < iters; i++ {
    58  		var doc bson.D
    59  		err := db.RunCommand(ctx, cmd).Decode(&doc)
    60  		if err != nil {
    61  			return err
    62  		}
    63  		// read the document and then throw it away to prevent
    64  		out, err := bson.Marshal(doc)
    65  		if err != nil {
    66  			return err
    67  		}
    68  		if len(out) == 0 {
    69  			return errors.New("output of command is empty")
    70  		}
    71  	}
    72  	tm.StopTimer()
    73  
    74  	return nil
    75  }
    76  
    77  func SingleFindOneByID(ctx context.Context, tm TimerManager, iters int) error {
    78  	ctx, cancel := context.WithCancel(ctx)
    79  	defer cancel()
    80  
    81  	db, err := getClientDB(ctx)
    82  	if err != nil {
    83  		return err
    84  	}
    85  
    86  	db = db.Client().Database("perftest")
    87  	if err = db.Drop(ctx); err != nil {
    88  		return err
    89  	}
    90  
    91  	doc, err := loadSourceDocument(getProjectRoot(), perfDataDir, singleAndMultiDataDir, tweetData)
    92  	if err != nil {
    93  		return err
    94  	}
    95  	coll := db.Collection("corpus")
    96  
    97  	for i := 0; i < iters; i++ {
    98  		idDoc := make(bson.D, 0, len(doc)+1)
    99  		idDoc = append(idDoc, bson.E{"_id", i})
   100  		idDoc = append(idDoc, doc...)
   101  		res, err := coll.InsertOne(ctx, idDoc)
   102  		if err != nil {
   103  			return err
   104  		}
   105  		if res.InsertedID == nil {
   106  			return errors.New("no inserted ID returned")
   107  		}
   108  	}
   109  
   110  	tm.ResetTimer()
   111  
   112  	for i := 0; i < iters; i++ {
   113  		var res bson.D
   114  		err := coll.FindOne(ctx, bson.D{{"_id", i}}).Decode(&res)
   115  		if err != nil {
   116  			return err
   117  		}
   118  	}
   119  
   120  	tm.StopTimer()
   121  
   122  	return db.Drop(ctx)
   123  }
   124  
   125  func singleInsertCase(ctx context.Context, tm TimerManager, iters int, data string) error {
   126  	ctx, cancel := context.WithCancel(ctx)
   127  	defer cancel()
   128  
   129  	db, err := getClientDB(ctx)
   130  	if err != nil {
   131  		return err
   132  	}
   133  	defer func() { _ = db.Client().Disconnect(ctx) }()
   134  
   135  	db = db.Client().Database("perftest")
   136  	if err = db.Drop(ctx); err != nil {
   137  		return err
   138  	}
   139  
   140  	doc, err := loadSourceDocument(getProjectRoot(), perfDataDir, singleAndMultiDataDir, data)
   141  	if err != nil {
   142  		return err
   143  	}
   144  
   145  	err = db.RunCommand(ctx, bson.D{{"create", "corpus"}}).Err()
   146  	if err != nil {
   147  		return err
   148  	}
   149  
   150  	coll := db.Collection("corpus")
   151  
   152  	tm.ResetTimer()
   153  
   154  	for i := 0; i < iters; i++ {
   155  		if _, err = coll.InsertOne(ctx, doc); err != nil {
   156  			return err
   157  		}
   158  	}
   159  
   160  	tm.StopTimer()
   161  
   162  	return db.Drop(ctx)
   163  }
   164  
   165  func SingleInsertSmallDocument(ctx context.Context, tm TimerManager, iters int) error {
   166  	return singleInsertCase(ctx, tm, iters, smallData)
   167  }
   168  
   169  func SingleInsertLargeDocument(ctx context.Context, tm TimerManager, iters int) error {
   170  	return singleInsertCase(ctx, tm, iters, largeData)
   171  }
   172
View as plain text