github.com/grailbio/bigslice@v0.0.0-20230519005545-30c4c12152ad/cmd/slicer/oom.go (about)

     1  // Copyright 2019 GRAIL, Inc. All rights reserved.
     2  // Use of this source code is governed by the Apache 2.0
     3  // license that can be found in the LICENSE file.
     4  
     5  package main
     6  
     7  import (
     8  	"context"
     9  	"flag"
    10  	"fmt"
    11  	"os"
    12  	"sync"
    13  
    14  	"github.com/grailbio/base/log"
    15  	"github.com/grailbio/base/stress/oom"
    16  	"github.com/grailbio/bigslice"
    17  	"github.com/grailbio/bigslice/exec"
    18  )
    19  
    20  var oomTest = bigslice.Func(func(size int) (slice bigslice.Slice) {
    21  	var mu sync.Mutex
    22  	slice = randomReader(100, 1000)
    23  	slice = bigslice.Map(slice, func(key string, values []int) string {
    24  		mu.Lock() // just one passes through per machine
    25  		if size != 0 {
    26  			oom.Do(size)
    27  		}
    28  		oom.Try()
    29  		panic("not reached")
    30  	})
    31  	return slice
    32  })
    33  
    34  func oomer(sess *exec.Session, args []string) error {
    35  	var (
    36  		flags = flag.NewFlagSet("oom", flag.ExitOnError)
    37  		size  = flags.Int("size", 0, "size of memory allocation; automatically determined if zero")
    38  	)
    39  	flags.Usage = func() {
    40  		fmt.Fprintln(os.Stderr, `usage: slicer oom [-size bytes]`)
    41  		flags.PrintDefaults()
    42  		os.Exit(2)
    43  	}
    44  	if err := flags.Parse(args); err != nil {
    45  		log.Fatal(err)
    46  	}
    47  
    48  	ctx := context.Background()
    49  	// Currently we only care about how OOMs are reported;
    50  	// we may try to recover from the in the future.
    51  	_, err := sess.Run(ctx, oomTest, *size)
    52  	return err
    53  }