github.com/NVIDIA/aistore@v1.3.23-0.20240517131212-7df6609be51d/cmd/aisnodeprofile/main.go (about)

     1  // Package main for the AIS node executable.
     2  /*
     3   * Copyright (c) 2018-2022, NVIDIA CORPORATION. All rights reserved.
     4   */
     5  package main
     6  
     7  import (
     8  	"flag"
     9  	"os"
    10  	"runtime"
    11  	"runtime/pprof"
    12  	"strconv"
    13  	"syscall"
    14  
    15  	"github.com/NVIDIA/aistore/ais"
    16  	"github.com/NVIDIA/aistore/cmn"
    17  	"github.com/NVIDIA/aistore/cmn/cos"
    18  	"github.com/NVIDIA/aistore/cmn/nlog"
    19  )
    20  
    21  var (
    22  	cpuProfile = flag.String("cpuprofile", "", "write cpu profile to `file`")
    23  	memProfile = flag.String("memprofile", "", "write memory profile to `file`")
    24  )
    25  
    26  // NOTE: these variables are set by ldflags
    27  var (
    28  	build     string
    29  	buildtime string
    30  )
    31  
    32  func main() {
    33  	os.Exit(run())
    34  }
    35  
    36  func run() int {
    37  	flag.Parse()
    38  
    39  	if s := *cpuProfile; s != "" {
    40  		*cpuProfile = s + "." + strconv.Itoa(syscall.Getpid())
    41  		f, err := os.Create(*cpuProfile)
    42  		if err != nil {
    43  			cos.ExitLogf("Couldn't create CPU profile: %v", err)
    44  		}
    45  		defer f.Close()
    46  		if err := pprof.StartCPUProfile(f); err != nil {
    47  			cos.ExitLogf("Couldn't start CPU profile: %v", err)
    48  		}
    49  		defer pprof.StopCPUProfile()
    50  	}
    51  
    52  	exitCode := ais.Run(cmn.VersionAIStore+"."+build, buildtime)
    53  	nlog.Flush(nlog.ActExit)
    54  
    55  	if s := *memProfile; s != "" {
    56  		*memProfile = s + "." + strconv.Itoa(syscall.Getpid())
    57  		f, err := os.Create(*memProfile)
    58  		if err != nil {
    59  			cos.ExitLogf("Couldn't create memory profile: %v", err)
    60  		}
    61  		defer f.Close()
    62  		runtime.GC() // get up-to-date statistics
    63  		if err := pprof.WriteHeapProfile(f); err != nil {
    64  			cos.ExitLogf("Couldn't write memory profile: %v", err)
    65  		}
    66  	}
    67  
    68  	return exitCode
    69  }