github.com/unigraph-dev/dgraph@v1.1.1-0.20200923154953-8b52b426f765/systest/posting-list-benchmark/main.go (about)

     1  /*
     2   * Copyright 2018 Dgraph Labs, Inc. and Contributors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package main
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"log"
    23  	"os"
    24  	"strings"
    25  	"time"
    26  
    27  	"github.com/dgraph-io/dgo"
    28  	"github.com/dgraph-io/dgo/protos/api"
    29  	"github.com/dgraph-io/dgraph/testutil"
    30  	"github.com/dgraph-io/dgraph/x"
    31  	"github.com/spf13/cobra"
    32  )
    33  
    34  type flagOptions struct {
    35  	addr            string
    36  	text            string
    37  	numMutations    uint64
    38  	mutationsPerTxn uint64
    39  	verbose         bool
    40  }
    41  
    42  var (
    43  	opt     flagOptions
    44  	rootCmd = &cobra.Command{
    45  		Use:   "posting-list-benchmark",
    46  		Short: "Benchmark to measure posting list performance.",
    47  		Long: `This benchmark adds a lot of triples with same object and predicate, creates a
    48  full-text index on that predicate, and measures the performance of mutations and
    49  transactions. The benchmark is intended to compare the performance of normal
    50  posting lists against the performance of multi-part posting lists.`,
    51  		Run: func(cmd *cobra.Command, args []string) {
    52  			runBenchmark()
    53  		},
    54  	}
    55  	defaultText = "One morning, when Gregor Samsa woke from troubled dreams, he found himself " +
    56  		"transformed in his bed into a horrible vermin. He lay on his armour-like back, and if " +
    57  		"he lifted his head a little he could see his brown belly, slightly domed and divided by " +
    58  		"arches into stiff sections. The bedding was hardly able to cover it and seemed ready to " +
    59  		"slide off any moment."
    60  )
    61  
    62  func initCmd() {
    63  	flags := rootCmd.Flags()
    64  	flags.StringVarP(&opt.addr, "addr", "a", "localhost:9180",
    65  		"Address of the Dgraph alpha to which updates will be sent.")
    66  	flags.StringVarP(&opt.text, "text", "t", defaultText,
    67  		"The English text to insert in each triple.")
    68  	flags.Uint64VarP(&opt.numMutations, "num_mutations", "", uint64(1e6),
    69  		"The total number of mutations that will be sent to the alpha.")
    70  	flags.Uint64VarP(&opt.mutationsPerTxn, "mutations_per_txn", "", uint64(1e3),
    71  		"The total number of mutations that will be sent in each transaction.")
    72  	flags.BoolVarP(&opt.verbose, "verbose", "v", false, "Toggle verbose output.")
    73  
    74  	if err := rootCmd.Execute(); err != nil {
    75  		fmt.Println(err)
    76  		os.Exit(1)
    77  	}
    78  }
    79  
    80  func runBenchmark() {
    81  	// Check the value of the flags are sane.
    82  	if opt.mutationsPerTxn > opt.numMutations {
    83  		log.Fatalf("Mutations per txn (%v) cannot be greater than the number of mutations (%v)",
    84  			opt.mutationsPerTxn, opt.numMutations)
    85  	}
    86  
    87  	dg := testutil.DgraphClientWithGroot(opt.addr)
    88  	var err error
    89  
    90  	// Drop all existing data.
    91  	for {
    92  		// keep retrying until we succeed or receive a non-retriable error
    93  		err = dg.Alter(context.Background(), &api.Operation{DropAll: true})
    94  		if err == nil || !strings.Contains(err.Error(), "Please retry") {
    95  			break
    96  		}
    97  		time.Sleep(time.Second)
    98  	}
    99  	x.Check(err)
   100  
   101  	// Initialize schema.
   102  	err = dg.Alter(context.Background(), &api.Operation{
   103  		Schema: "text: string @index(fulltext) @lang .",
   104  	})
   105  	x.Check(err)
   106  
   107  	var triples []string
   108  	for i := uint64(0); i < opt.numMutations; i++ {
   109  		uid := fmt.Sprintf("_:uid%d", i)
   110  		triple := fmt.Sprintf("%s <text> \"%s\"@en .", uid, opt.text)
   111  		triples = append(triples, triple)
   112  
   113  		if i > 0 && i%opt.mutationsPerTxn == 0 {
   114  			commitTriples(dg, triples)
   115  			triples = nil
   116  		}
   117  	}
   118  
   119  	// Commit last transaction in case it has not still been done.
   120  	commitTriples(dg, triples)
   121  }
   122  
   123  func commitTriples(dg *dgo.Dgraph, triples []string) {
   124  	for {
   125  		txn := dg.NewTxn()
   126  		mu := api.Mutation{
   127  			SetNquads: []byte(strings.Join(triples, "\n")),
   128  			CommitNow: true,
   129  		}
   130  		_, err := txn.Mutate(context.Background(), &mu)
   131  
   132  		if err == nil {
   133  			break
   134  		}
   135  
   136  		// Retry in case the transaction has been aborted.
   137  		if strings.Contains(err.Error(), "Transaction has been aborted") {
   138  			continue
   139  		}
   140  
   141  		x.Check(err)
   142  	}
   143  }
   144  
   145  func main() {
   146  	initCmd()
   147  }