diff --git a/cmd/cmd_run.go b/cmd/cmd_run.go index fada561..61c8a23 100644 --- a/cmd/cmd_run.go +++ b/cmd/cmd_run.go @@ -28,6 +28,7 @@ import ( runesdatagateway "github.com/gaze-network/indexer-network/modules/runes/datagateway" runespostgres "github.com/gaze-network/indexer-network/modules/runes/repository/postgres" runesusecase "github.com/gaze-network/indexer-network/modules/runes/usecase" + "github.com/gaze-network/indexer-network/pkg/automaxprocs" "github.com/gaze-network/indexer-network/pkg/errorhandler" "github.com/gaze-network/indexer-network/pkg/logger" "github.com/gaze-network/indexer-network/pkg/logger/slogx" @@ -57,6 +58,9 @@ func NewRunCommand() *cobra.Command { Use: "run", Short: "Start indexer-network service", RunE: func(cmd *cobra.Command, args []string) error { + if err := automaxprocs.Init(); err != nil { + logger.Error("Failed to set GOMAXPROCS", slogx.Error(err)) + } return runHandler(opts, cmd, args) }, } diff --git a/main.go b/main.go index 61f0e17..6906216 100644 --- a/main.go +++ b/main.go @@ -7,7 +7,6 @@ import ( "syscall" "github.com/gaze-network/indexer-network/cmd" - _ "go.uber.org/automaxprocs" ) func main() { diff --git a/pkg/automaxprocs/automaxprocs.go b/pkg/automaxprocs/automaxprocs.go new file mode 100644 index 0000000..c82a410 --- /dev/null +++ b/pkg/automaxprocs/automaxprocs.go @@ -0,0 +1,99 @@ +package automaxprocs + +import ( + "context" + "fmt" + "log/slog" + "os" + "runtime" + + "github.com/Cleverse/go-utilities/utils" + "github.com/cockroachdb/errors" + "github.com/gaze-network/indexer-network/pkg/logger" + "github.com/gaze-network/indexer-network/pkg/logger/slogx" + "go.uber.org/automaxprocs/maxprocs" +) + +var ( + // undo is the undo function returned by maxprocs.Set + undo func() + + // autoMaxProcs is the value of GOMAXPROCS set by automaxprocs. + // will be -1 if `automaxprocs` is not initialized. + autoMaxProcs = -1 + + // initialMaxProcs is the initial value of GOMAXPROCS. + initialMaxProcs = Current() +) + +func Init() error { + logger := logger.With( + slogx.String("package", "automaxprocs"), + slogx.String("event", "set_gomaxprocs"), + slogx.Int("prev_maxprocs", initialMaxProcs), + ) + + // Create a logger function for `maxprocs.Set`. + setMaxProcLogger := func(format string, v ...any) { + fields := make([]slog.Attr, 0, 1) + + // `maxprocs.Set` will always pass current GOMAXPROCS value to logger. + // except when calling `undo` function, it will not pass any value. + if val, ok := utils.Optional(v); ok { + // if `GOMAXPROCS` environment variable is set, then `automaxprocs` will honor it. + if _, exists := os.LookupEnv("GOMAXPROCS"); exists { + val = Current() + } + + // add logging field for `set_maxprocs` value if it's present in integer value. + if setmaxprocs, ok := val.(int); ok { + fields = append(fields, slogx.Int("set_maxprocs", setmaxprocs)) + } + } + + logger.LogAttrs(context.Background(), slog.LevelInfo, fmt.Sprintf(format, v...), fields...) + } + + // Set GOMAXPROCS to match the Linux container CPU quota (if any), returning + // any error encountered and an undo function. + // + // Set is a no-op on non-Linux systems and in Linux environments without a + // configured CPU quota. + revert, err := maxprocs.Set(maxprocs.Logger(setMaxProcLogger), maxprocs.Min(1)) + if err != nil { + return errors.WithStack(err) + } + + // set the result of `maxprocs.Set` to global variable. + autoMaxProcs = Current() + undo = revert + return nil +} + +// Undo restores GOMAXPROCS to its previous value. +// or revert to initial value if `automaxprocs` is not initialized. +// +// returns the current GOMAXPROCS value. +func Undo() int { + if undo != nil { + undo() + return Current() + } + + runtime.GOMAXPROCS(initialMaxProcs) + return initialMaxProcs +} + +// Current returns the current value of GOMAXPROCS. +func Current() int { + return runtime.GOMAXPROCS(0) +} + +// Value returns the value of GOMAXPROCS set by automaxprocs. +// returns -1 if `automaxprocs` is not initialized. +func Value() int { + if autoMaxProcs <= 0 { + return -1 + } + return autoMaxProcs +}