Skip to content

Commit

Permalink
roachprod: load balancer expanders
Browse files Browse the repository at this point in the history
Add "L" to the node expansion in `roachprod` (e.g., `test-cluster:L` vs.
`test-cluster:1`). This allows users to pass "L" instead of a node range /
number. It will return the load balancer IP or postgres URL (depending on the
expander used) and will enable `workload` and other utilities to target the load
balancer IP rather than a specific node.

Epic: None
Release Note: None
  • Loading branch information
herkolategan committed Apr 24, 2024
1 parent 0f56ecd commit 4edacf9
Showing 1 changed file with 41 additions and 14 deletions.
55 changes: 41 additions & 14 deletions pkg/roachprod/install/expander.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ import (
)

var parameterRe = regexp.MustCompile(`{[^{}]*}`)
var pgURLRe = regexp.MustCompile(`{pgurl(:[-,0-9]+)?(:[a-z0-9\-]+)?(:[0-9]+)?}`)
var pgHostRe = regexp.MustCompile(`{pghost(:[-,0-9]+)?}`)
var pgURLRe = regexp.MustCompile(`{pgurl(:[-,0-9]+|:L)?(:[a-z0-9\-]+)?(:[0-9]+)?}`)
var pgHostRe = regexp.MustCompile(`{pghost(:[-,0-9]+|:L)?(:[a-z0-9\-]+)?(:[0-9]+)?}`)
var pgPortRe = regexp.MustCompile(`{pgport(:[-,0-9]+)?(:[a-z0-9\-]+)?(:[0-9]+)?}`)
var uiPortRe = regexp.MustCompile(`{uiport(:[-,0-9]+)}`)
var storeDirRe = regexp.MustCompile(`{store-dir(:[0-9]+)?}`)
Expand Down Expand Up @@ -157,14 +157,20 @@ func (e *expander) maybeExpandPgURL(
if err != nil {
return "", false, err
}
if e.pgURLs[virtualClusterName] == nil {
e.pgURLs[virtualClusterName], err = c.pgurls(ctx, l, allNodes(len(c.VMs)), virtualClusterName, sqlInstance)
if err != nil {
return "", false, err
switch m[1] {
case ":L":
url, err := c.loadBalancerURL(ctx, l, virtualClusterName, sqlInstance)
return url, url != "", err
default:
if e.pgURLs[virtualClusterName] == nil {
e.pgURLs[virtualClusterName], err = c.pgurls(ctx, l, allNodes(len(c.VMs)), virtualClusterName, sqlInstance)
if err != nil {
return "", false, err
}
}
s, err = e.maybeExpandMap(c, e.pgURLs[virtualClusterName], m[1])
return s, err == nil, err
}
s, err = e.maybeExpandMap(c, e.pgURLs[virtualClusterName], m[1])
return s, err == nil, err
}

// maybeExpandPgHost is an expanderFunc for {pghost:<nodeSpec>}
Expand All @@ -175,17 +181,38 @@ func (e *expander) maybeExpandPgHost(
if m == nil {
return s, false, nil
}
virtualClusterName, sqlInstance, err := extractVirtualClusterInfo(m[2:])
if err != nil {
return "", false, err
}

if e.pgHosts == nil {
var err error
e.pgHosts, err = c.pghosts(ctx, l, allNodes(len(c.VMs)))
switch m[1] {
case ":L":
services, err := c.DiscoverServices(ctx, virtualClusterName, ServiceTypeSQL, ServiceInstancePredicate(sqlInstance))
if err != nil {
return "", false, err
}
for _, svc := range services {
if svc.VirtualClusterName == virtualClusterName && svc.Instance == sqlInstance {
addr, err := c.FindLoadBalancer(l, svc.Port)
if err != nil {
return "", false, err
}
return addr.IP, true, nil
}
}
return "", false, err
default:
if e.pgHosts == nil {
var err error
e.pgHosts, err = c.pghosts(ctx, l, allNodes(len(c.VMs)))
if err != nil {
return "", false, err
}
}
s, err := e.maybeExpandMap(c, e.pgHosts, m[1])
return s, err == nil, err
}

s, err := e.maybeExpandMap(c, e.pgHosts, m[1])
return s, err == nil, err
}

// maybeExpandPgURL is an expanderFunc for {pgport:<nodeSpec>}
Expand Down

0 comments on commit 4edacf9

Please sign in to comment.