main.go

// SPDX-FileCopyrightText: 2021 Dinko Korunic
// SPDX-License-Identifier: MIT

package main

import (
	"fmt"
	"net/netip"
	"os"
	"runtime"
	"runtime/pprof"
	"sync"

	"github.com/KimMachineGun/automemlimit/memlimit"
)

const (
	mapSize      = 4096
	subMapSize   = 8
	hostChanSize = 2048
	maxMemRatio  = 0.9
)

var (
	GitTag    = ""
	GitCommit = ""
	GitDirty  = ""
	BuildTime = ""
)

func main() {
	_, _ = memlimit.SetGoMemLimitWithOpts(
		memlimit.WithRatio(maxMemRatio),
		memlimit.WithProvider(
			memlimit.ApplyFallback(
				memlimit.FromCgroup,
				memlimit.FromSystem,
			),
		),
	)

	zones, server, cidrList := parseFlags()

	if *cpuProfile != "" {
		f, err := os.Create(*cpuProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error creating CPU profile: %v\n", err)
		} else {
			defer f.Close()

			if err := pprof.StartCPUProfile(f); err != nil {
				fmt.Fprintf(os.Stderr, "Error starting CPU profile: %v\n", err)
			}
			defer pprof.StopCPUProfile()
		}
	}

	if *memProfile != "" {
		f, err := os.Create(*memProfile)
		if err != nil {
			fmt.Fprintf(os.Stderr, "Error trying to create memory profile: %v\n", err)
		} else {
			defer f.Close()

			defer func() {
				runtime.GC()

				if err := pprof.WriteHeapProfile(f); err != nil {
					fmt.Fprintf(os.Stderr, "Error writing memory profile: %v\n", err)
				}
			}()
		}
	}

	_ = setNofile()

	ranger, doCIDR := rangerInit(cidrList)
	hostChan := make(chan HostEntry, hostChanSize)

	entries := make(HostMap, mapSize)
	keys := make([]netip.Addr, 0, mapSize)

	var wgMon, wgWrk sync.WaitGroup

	// sole owner of entries/keys; lock-free via channel
	wgMon.Go(func() {
		writeHostEntries(hostChan, &keys, entries)
	})

	semAXFR := make(chan struct{}, *maxTransfers)

	for _, zone := range zones {
		if server == "" {
			wgWrk.Go(func() {
				processLocalZone(zone, doCIDR, ranger, hostChan)
			})
		} else {
			wgWrk.Add(1)
			semAXFR <- struct{}{}

			go func() {
				defer wgWrk.Done()

				processRemoteZone(zone, server, doCIDR, ranger, hostChan)
				<-semAXFR
			}()
		}
	}

	wgWrk.Wait()
	close(hostChan)
	wgMon.Wait()

	displayHostEntries(keys, entries)
}