github.com/onflow/atree@v0.6.0/cmd/stress/main.go (about)

     1  /*
     2   * Atree - Scalable Arrays and Ordered Maps
     3   *
     4   * Copyright 2021 Dapper Labs, Inc.
     5   *
     6   * Licensed under the Apache License, Version 2.0 (the "License");
     7   * you may not use this file except in compliance with the License.
     8   * You may obtain a copy of the License at
     9   *
    10   *   http://www.apache.org/licenses/LICENSE-2.0
    11   *
    12   * Unless required by applicable law or agreed to in writing, software
    13   * distributed under the License is distributed on an "AS IS" BASIS,
    14   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    15   * See the License for the specific language governing permissions and
    16   * limitations under the License.
    17   */
    18  
    19  package main
    20  
    21  import (
    22  	"flag"
    23  	"fmt"
    24  	"io"
    25  	"os"
    26  	"os/signal"
    27  	"strconv"
    28  	"strings"
    29  	"syscall"
    30  	"time"
    31  
    32  	"github.com/onflow/atree"
    33  
    34  	"github.com/fxamacker/cbor/v2"
    35  )
    36  
    37  const maxStatusLength = 128
    38  
    39  type Status interface {
    40  	Write()
    41  }
    42  
    43  func writeStatus(status string) {
    44  	// Clear old status
    45  	s := fmt.Sprintf("\r%s\r", strings.Repeat(" ", maxStatusLength))
    46  	_, _ = io.WriteString(os.Stdout, s)
    47  
    48  	// Write new status
    49  	_, _ = io.WriteString(os.Stdout, status)
    50  }
    51  
    52  func updateStatus(sigc <-chan os.Signal, status Status) {
    53  
    54  	status.Write()
    55  
    56  	ticker := time.NewTicker(3 * time.Second)
    57  
    58  	for {
    59  		select {
    60  		case <-ticker.C:
    61  			status.Write()
    62  
    63  		case <-sigc:
    64  			status.Write()
    65  			fmt.Fprintf(os.Stdout, "\n")
    66  
    67  			ticker.Stop()
    68  			os.Exit(1)
    69  		}
    70  	}
    71  }
    72  
    73  func main() {
    74  
    75  	var typ string
    76  	var maxLength uint64
    77  	var seedHex string
    78  	var minHeapAllocMiB, maxHeapAllocMiB uint64
    79  
    80  	flag.StringVar(&typ, "type", "array", "array or map")
    81  	flag.Uint64Var(&maxLength, "maxlen", 10_000, "max number of elements")
    82  	flag.StringVar(&seedHex, "seed", "", "seed for prng in hex (default is Unix time)")
    83  	flag.Uint64Var(&minHeapAllocMiB, "minheap", 1000, "min HeapAlloc in MiB to stop extra removal of elements")
    84  	flag.Uint64Var(&maxHeapAllocMiB, "maxheap", 2000, "max HeapAlloc in MiB to trigger extra removal of elements")
    85  
    86  	flag.Parse()
    87  
    88  	var seed int64
    89  	if len(seedHex) != 0 {
    90  		var err error
    91  		seed, err = strconv.ParseInt(strings.ReplaceAll(seedHex, "0x", ""), 16, 64)
    92  		if err != nil {
    93  			panic("Failed to parse seed flag (hex string)")
    94  		}
    95  	}
    96  
    97  	r = newRand(seed)
    98  
    99  	typ = strings.ToLower(typ)
   100  
   101  	if typ != "array" && typ != "map" {
   102  		fmt.Fprintf(os.Stderr, "Please specify type as either \"array\" or \"map\"")
   103  		return
   104  	}
   105  
   106  	sigc := make(chan os.Signal, 1)
   107  	signal.Notify(sigc, os.Interrupt, syscall.SIGTERM)
   108  
   109  	// Create storage
   110  	encMode, err := cbor.EncOptions{}.EncMode()
   111  	if err != nil {
   112  		fmt.Fprintf(os.Stderr, "Failed to create CBOR encoding mode: %s\n", err)
   113  		return
   114  	}
   115  
   116  	decMode, err := cbor.DecOptions{}.DecMode()
   117  	if err != nil {
   118  		fmt.Fprintf(os.Stderr, "Failed to create CBOR decoding mode: %s\n", err)
   119  		return
   120  	}
   121  
   122  	baseStorage := NewInMemBaseStorage()
   123  
   124  	storage := atree.NewPersistentSlabStorage(
   125  		baseStorage,
   126  		encMode,
   127  		decMode,
   128  		decodeStorable,
   129  		decodeTypeInfo,
   130  	)
   131  
   132  	typeInfo := testTypeInfo{value: 123}
   133  
   134  	address := atree.Address{1, 2, 3, 4, 5, 6, 7, 8}
   135  
   136  	switch typ {
   137  
   138  	case "array":
   139  		fmt.Printf("Starting array stress test, minMapHeapAlloc = %d MiB, maxMapHeapAlloc = %d MiB\n", minHeapAllocMiB, maxHeapAllocMiB)
   140  
   141  		status := newArrayStatus()
   142  
   143  		go updateStatus(sigc, status)
   144  
   145  		testArray(storage, address, typeInfo, maxLength, status, minHeapAllocMiB, maxHeapAllocMiB)
   146  
   147  	case "map":
   148  		fmt.Printf("Starting map stress test, minMapHeapAlloc = %d MiB, maxMapHeapAlloc = %d MiB\n", minHeapAllocMiB, maxHeapAllocMiB)
   149  
   150  		status := newMapStatus()
   151  
   152  		go updateStatus(sigc, status)
   153  
   154  		testMap(storage, address, typeInfo, maxLength, status, minHeapAllocMiB, maxHeapAllocMiB)
   155  	}
   156  
   157  }