Skip to content

Commit

Permalink
fix: health check
Browse files Browse the repository at this point in the history
  • Loading branch information
@dennis-tra
dennis-tra committed Dec 3, 2024
1 parent 1222727 commit 74babba
Show file tree
Hide file tree
Showing 2 changed files with 83 additions and 61 deletions.
14 changes: 10 additions & 4 deletions cmd/ants/health.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,16 +3,22 @@ package main
import (
"fmt"
"net/http"
"os"

"github.com/urfave/cli/v2"
)

var healthConfig = struct {
MetricsHost string
MetricsPort int
}{
MetricsHost: "127.0.0.1",
MetricsPort: 5999, // one below the FirstPort to not accidentally override it
}

func HealthCheck(c *cli.Context) error {
endpoint := fmt.Sprintf(
"http://%s:%s/health",
os.Getenv("METRICS_HOST"),
os.Getenv("METRICS_PORT"),
"http://%s:%d/metrics",
healthConfig.MetricsHost, healthConfig.MetricsPort,
)
req, err := http.NewRequestWithContext(c.Context, http.MethodGet, endpoint, nil)
if err != nil {
Expand Down
130 changes: 73 additions & 57 deletions cmd/ants/main.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import (

var logger = logging.Logger("ants-queen")

var rootConfig = struct {
var queenConfig = struct {
MetricsHost string
MetricsPort int
ClickhouseAddress string
Expand Down Expand Up @@ -79,134 +79,134 @@ func main() {
Name: "metrics.host",
Usage: "On which host to expose the metrics",
EnvVars: []string{"ANTS_METRICS_HOST"},
Destination: &rootConfig.MetricsHost,
Value: rootConfig.MetricsHost,
Destination: &queenConfig.MetricsHost,
Value: queenConfig.MetricsHost,
},
&cli.IntFlag{
Name: "metrics.port",
Usage: "On which port to expose the metrics",
EnvVars: []string{"ANTS_METRICS_PORT"},
Destination: &rootConfig.MetricsPort,
Value: rootConfig.MetricsPort,
Destination: &queenConfig.MetricsPort,
Value: queenConfig.MetricsPort,
},
&cli.StringFlag{
Name: "clickhouse.address",
Usage: "ClickHouse address containing the host and port, 127.0.0.1:9000",
EnvVars: []string{"ANTS_CLICKHOUSE_ADDRESS"},
Destination: &rootConfig.ClickhouseAddress,
Value: rootConfig.ClickhouseAddress,
Destination: &queenConfig.ClickhouseAddress,
Value: queenConfig.ClickhouseAddress,
},
&cli.StringFlag{
Name: "clickhouse.database",
Usage: "The ClickHouse database where ants requests will be recorded",
EnvVars: []string{"ANTS_CLICKHOUSE_DATABASE"},
Destination: &rootConfig.ClickhouseDatabase,
Value: rootConfig.ClickhouseDatabase,
Destination: &queenConfig.ClickhouseDatabase,
Value: queenConfig.ClickhouseDatabase,
},
&cli.StringFlag{
Name: "clickhouse.username",
Usage: "The ClickHouse user that has the prerequisite privileges to record the requests",
EnvVars: []string{"ANTS_CLICKHOUSE_USERNAME"},
Destination: &rootConfig.ClickhouseUsername,
Value: rootConfig.ClickhouseUsername,
Destination: &queenConfig.ClickhouseUsername,
Value: queenConfig.ClickhouseUsername,
},
&cli.StringFlag{
Name: "clickhouse.password",
Usage: "The password for the ClickHouse user",
EnvVars: []string{"ANTS_CLICKHOUSE_PASSWORD"},
Destination: &rootConfig.ClickhousePassword,
Value: rootConfig.ClickhousePassword,
Destination: &queenConfig.ClickhousePassword,
Value: queenConfig.ClickhousePassword,
},
&cli.BoolFlag{
Name: "clickhouse.ssl",
Usage: "Whether to use SSL for the ClickHouse connection",
EnvVars: []string{"ANTS_CLICKHOUSE_SSL"},
Destination: &rootConfig.ClickhouseSSL,
Value: rootConfig.ClickhouseSSL,
Destination: &queenConfig.ClickhouseSSL,
Value: queenConfig.ClickhouseSSL,
},
&cli.StringFlag{
Name: "nebula.connstring",
Usage: "The connection string for the Postgres Nebula database",
EnvVars: []string{"ANTS_NEBULA_CONNSTRING"},
Destination: &rootConfig.NebulaDBConnString,
Value: rootConfig.NebulaDBConnString,
Destination: &queenConfig.NebulaDBConnString,
Value: queenConfig.NebulaDBConnString,
},
&cli.IntFlag{
Name: "batch.size",
Usage: "The number of ants to request to store at a time",
EnvVars: []string{"ANTS_BATCH_SIZE"},
Destination: &rootConfig.BatchSize,
Value: rootConfig.BatchSize,
Destination: &queenConfig.BatchSize,
Value: queenConfig.BatchSize,
},
&cli.DurationFlag{
Name: "batch.time",
Usage: "The time to wait between batches",
EnvVars: []string{"ANTS_BATCH_TIME"},
Destination: &rootConfig.BatchTime,
Value: rootConfig.BatchTime,
Destination: &queenConfig.BatchTime,
Value: queenConfig.BatchTime,
},
&cli.DurationFlag{
Name: "crawl.interval",
Usage: "The time between two crawls",
EnvVars: []string{"ANTS_CRAWL_INTERVAL"},
Destination: &rootConfig.CrawlInterval,
Value: rootConfig.CrawlInterval,
Destination: &queenConfig.CrawlInterval,
Value: queenConfig.CrawlInterval,
},
&cli.IntFlag{
Name: "cache.size",
Usage: "How many agent versions and protocols should be cached in memory",
EnvVars: []string{"ANTS_CACHE_SIZE"},
Destination: &rootConfig.CacheSize,
Value: rootConfig.CacheSize,
Destination: &queenConfig.CacheSize,
Value: queenConfig.CacheSize,
},
&cli.PathFlag{
Name: "key.path",
Usage: "The path to the data store containing the keys",
EnvVars: []string{"ANTS_KEY_PATH"},
Destination: &rootConfig.KeyDBPath,
Value: rootConfig.KeyDBPath,
Destination: &queenConfig.KeyDBPath,
Value: queenConfig.KeyDBPath,
},
&cli.IntFlag{
Name: "first_port",
Usage: "First port ants can listen on",
EnvVars: []string{"ANTS_FIRST_PORT"},
Destination: &rootConfig.FirstPort,
Value: rootConfig.FirstPort,
Destination: &queenConfig.FirstPort,
Value: queenConfig.FirstPort,
},
&cli.IntFlag{
Name: "num_ports",
Usage: "Number of ports ants can listen on",
EnvVars: []string{"ANTS_NUM_PORTS"},
Destination: &rootConfig.NumPorts,
Value: rootConfig.NumPorts,
Destination: &queenConfig.NumPorts,
Value: queenConfig.NumPorts,
},
&cli.BoolFlag{
Name: "upnp",
Usage: "Enable UPnP",
EnvVars: []string{"ANTS_UPNP"},
Destination: &rootConfig.UPnp,
Value: rootConfig.UPnp,
Destination: &queenConfig.UPnp,
Value: queenConfig.UPnp,
},
&cli.IntFlag{
Name: "bucket.size",
Usage: "The bucket size for the ants DHT",
EnvVars: []string{"ANTS_BUCKET_SIZE"},
Destination: &rootConfig.BucketSize,
Value: rootConfig.BucketSize,
Destination: &queenConfig.BucketSize,
Value: queenConfig.BucketSize,
},
&cli.StringFlag{
Name: "user.agent",
Usage: "The user agent to use for the ants hosts",
EnvVars: []string{"ANTS_USER_AGENT"},
Destination: &rootConfig.UserAgent,
Value: rootConfig.UserAgent,
Destination: &queenConfig.UserAgent,
Value: queenConfig.UserAgent,
},
&cli.StringFlag{
Name: "queen.id",
Usage: "The ID for the queen that's orchestrating the ants",
EnvVars: []string{"ANTS_QUEEN_ID"},
Destination: &rootConfig.QueenID,
Value: rootConfig.QueenID,
Destination: &queenConfig.QueenID,
Value: queenConfig.QueenID,
DefaultText: "generated",
},
},
Expand All @@ -216,6 +216,22 @@ func main() {
Name: "health",
Usage: "Checks the health of the service",
Action: HealthCheck,
Flags: []cli.Flag{
&cli.StringFlag{
Name: "metrics.host",
Usage: "On which host to expose the metrics",
EnvVars: []string{"ANTS_METRICS_HOST"},
Destination: &healthConfig.MetricsHost,
Value: healthConfig.MetricsHost,
},
&cli.IntFlag{
Name: "metrics.port",
Usage: "On which port to expose the metrics",
EnvVars: []string{"ANTS_METRICS_PORT"},
Destination: &healthConfig.MetricsPort,
Value: healthConfig.MetricsPort,
},
},
},
},
}
Expand Down Expand Up @@ -247,16 +263,16 @@ func runQueenCommand(c *cli.Context) error {
return fmt.Errorf("init telemetry: %w", err)
}

logger.Debugln("Starting metrics server", "host", rootConfig.MetricsHost, "port", rootConfig.MetricsPort)
go metrics.ListenAndServe(rootConfig.MetricsHost, rootConfig.MetricsPort)
logger.Debugln("Starting metrics server", "host", queenConfig.MetricsHost, "port", queenConfig.MetricsPort)
go metrics.ListenAndServe(queenConfig.MetricsHost, queenConfig.MetricsPort)

// initializing a new clickhouse client
client, err := db.NewClient(
rootConfig.ClickhouseAddress,
rootConfig.ClickhouseDatabase,
rootConfig.ClickhouseUsername,
rootConfig.ClickhousePassword,
rootConfig.ClickhouseSSL,
queenConfig.ClickhouseAddress,
queenConfig.ClickhouseDatabase,
queenConfig.ClickhouseUsername,
queenConfig.ClickhousePassword,
queenConfig.ClickhouseSSL,
telemetry,
)
if err != nil {
Expand All @@ -271,17 +287,17 @@ func runQueenCommand(c *cli.Context) error {
}

queenCfg := &ants.QueenConfig{
KeysDBPath: rootConfig.KeyDBPath,
NPorts: rootConfig.NumPorts,
FirstPort: rootConfig.FirstPort,
UPnP: rootConfig.UPnp,
BatchSize: rootConfig.BatchSize,
BatchTime: rootConfig.BatchTime,
CrawlInterval: rootConfig.CrawlInterval,
CacheSize: rootConfig.CacheSize,
NebulaDBConnString: rootConfig.NebulaDBConnString,
BucketSize: rootConfig.BucketSize,
UserAgent: rootConfig.UserAgent,
KeysDBPath: queenConfig.KeyDBPath,
NPorts: queenConfig.NumPorts,
FirstPort: queenConfig.FirstPort,
UPnP: queenConfig.UPnp,
BatchSize: queenConfig.BatchSize,
BatchTime: queenConfig.BatchTime,
CrawlInterval: queenConfig.CrawlInterval,
CacheSize: queenConfig.CacheSize,
NebulaDBConnString: queenConfig.NebulaDBConnString,
BucketSize: queenConfig.BucketSize,
UserAgent: queenConfig.UserAgent,
Telemetry: telemetry,
}

Expand Down

0 comments on commit 74babba

Please sign in to comment.