Skip to content

Commit 5542f56

Browse files
afiuneafiune
authored
✨ cache aws connection and upstream client (#5461)
This is an attempt to reduce the load of the AWS provider. We are going to cache the upstream client which uses ranger under the hood and creates goroutines to fulfill the requests. This is not a problem for a few runtimes, but when we do thousands of them, we really don't need to create that many clients. --------- Signed-off-by: Salim Afiune Maya <afiune@mondoo.com>
1 parent 7933370 commit 5542f56

File tree

6 files changed

+137
-68
lines changed

6 files changed

+137
-68
lines changed

‎providers/aws/connection/awsec2ebsconn/provider.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,11 @@ func NewAwsEbsConnection(id uint32, conf *inventory.Config, asset *inventory.Ass
174174
return c, nil
175175
}
176176

177+
// no-op since ebs connection doesn't do any verification
178+
func (c *AwsEbsConnection) Hash() uint64 { return 0 }
179+
func (c *AwsEbsConnection) Verify() (string, error) { return "no-op", nil }
180+
func (c *AwsEbsConnection) SetAccountId(_ string) {}
181+
177182
func (c *AwsEbsConnection) FileInfo(path string) (shared.FileInfoDetails, error) {
178183
return shared.FileInfoDetails{}, errors.New("FileInfo not implemented")
179184
}

‎providers/aws/connection/connection.go‎

Lines changed: 70 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ import (
88
"encoding/json"
99
"errors"
1010
"io"
11+
"maps"
1112
"net/http"
1213
"slices"
1314
"strconv"
@@ -22,6 +23,7 @@ import (
2223
"github.com/aws/aws-sdk-go-v2/service/ec2"
2324
"github.com/aws/aws-sdk-go-v2/service/sts"
2425
"github.com/hashicorp/go-retryablehttp"
26+
"github.com/mitchellh/hashstructure/v2"
2527
"github.com/rs/zerolog/log"
2628
"github.com/spf13/afero"
2729
"go.mondoo.com/cnquery/v11/logger/zerologadapter"
@@ -33,17 +35,22 @@ import (
3335

3436
type AwsConnection struct {
3537
plugin.Connection
36-
Conf *inventory.Config
37-
asset *inventory.Asset
38-
cfg aws.Config
39-
accountId string
40-
clientcache ClientsCache
41-
awsConfigOptions []func(*config.LoadOptions) error
42-
profile string
43-
PlatformOverride string
44-
connectionOptions map[string]string
45-
Filters DiscoveryFilters
46-
scope string
38+
Conf *inventory.Config
39+
asset *inventory.Asset
40+
cfg aws.Config
41+
accountId string
42+
clientcache ClientsCache
43+
awsConfigOptions []func(*config.LoadOptions) error
44+
PlatformOverride string
45+
Filters DiscoveryFilters
46+
47+
opts awsConnectionOptions
48+
}
49+
50+
type awsConnectionOptions struct {
51+
scope string
52+
profile string
53+
options map[string]string
4754
}
4855

4956
type DiscoveryFilters struct {
@@ -100,7 +107,11 @@ func NewAwsConnection(id uint32, asset *inventory.Asset, conf *inventory.Config)
100107
awsConfigOptions: []func(*config.LoadOptions) error{},
101108
Filters: EmptyDiscoveryFilters(),
102109
}
103-
opts := parseFlagsForConnectionOptions(asset.Options, conf.Options, conf.GetCredentials())
110+
111+
// merge the options to make sure we don't miss anything
112+
maps.Copy(asset.Options, conf.Options)
113+
114+
opts := parseFlagsForConnectionOptions(asset.Options, conf.GetCredentials())
104115
for _, opt := range opts {
105116
opt(c)
106117
}
@@ -118,34 +129,60 @@ func NewAwsConnection(id uint32, asset *inventory.Asset, conf *inventory.Config)
118129
log.Info().Msg("no AWS region found, using us-east-1")
119130
cfg.Region = "us-east-1" // in case the user has no region set, default to us-east-1
120131
}
121-
// gather information about the aws account
122-
cfgCopy := cfg.Copy()
123-
identity, err := CheckIam(cfgCopy)
124-
if err != nil {
125-
log.Debug().Err(err).Msg("could not gather details of AWS account")
126-
// try with govcloud region
127-
cfgCopy.Region = "us-gov-west-1"
128-
identity, err = CheckIam(cfgCopy)
129-
if err != nil {
130-
log.Debug().Err(err).Msg("could not gather details of AWS account")
131-
return nil, err
132-
}
133-
}
134132

135133
c.Connection = plugin.NewConnection(id, asset)
136134
c.Conf = conf
137135
c.asset = asset
138136
c.cfg = cfg
139-
c.accountId = *identity.Account
140-
c.profile = asset.Options["profile"]
141-
c.scope = asset.Options["scope"]
142-
c.connectionOptions = asset.Options
137+
c.opts.profile = asset.Options["profile"]
138+
c.opts.scope = asset.Options["scope"]
139+
c.opts.options = asset.Options
143140
if conf.Discover != nil {
144141
c.Filters = parseOptsToFilters(conf.Discover.Filter)
145142
}
146143
return c, nil
147144
}
148145

146+
func (c *AwsConnection) Hash() uint64 {
147+
// generate hash of the config options used to to initialize this connection,
148+
// we use this to avoid verifying a client with the same options more than once
149+
hash, err := hashstructure.Hash(c.opts, hashstructure.FormatV2, nil)
150+
if err != nil {
151+
log.Error().Err(err).Msg("unable to hash connection")
152+
}
153+
return hash
154+
}
155+
156+
func (c *AwsConnection) Verify() (string, error) {
157+
identity, err := CheckIam(c.cfg.Copy())
158+
if err != nil {
159+
log.Debug().Err(err).Msg("could not gather details of AWS account")
160+
// try with govcloud region, store error to return it if this last option does not work
161+
err1 := err
162+
cfgCopy := c.cfg.Copy()
163+
cfgCopy.Region = "us-gov-west-1"
164+
identity, err = CheckIam(cfgCopy)
165+
if err != nil {
166+
return "", err1
167+
}
168+
}
169+
account := ""
170+
if identity.Account != nil {
171+
account = *identity.Account
172+
}
173+
174+
return account, nil
175+
}
176+
177+
func (c *AwsConnection) SetAccountId(id string) {
178+
if id != "" {
179+
c.accountId = id
180+
}
181+
}
182+
func (p *AwsConnection) AccountId() string {
183+
return p.accountId
184+
}
185+
149186
func parseOptsToFilters(opts map[string]string) DiscoveryFilters {
150187
d := EmptyDiscoveryFilters()
151188
for k, v := range opts {
@@ -186,15 +223,7 @@ func parseOptsToFilters(opts map[string]string) DiscoveryFilters {
186223
return d
187224
}
188225

189-
func parseFlagsForConnectionOptions(m1 map[string]string, m2 map[string]string, creds []*vault.Credential) []ConnectionOption {
190-
// merge the options to make sure we don't miss anything
191-
m := m1
192-
if m == nil {
193-
m = make(map[string]string)
194-
}
195-
for k, v := range m2 {
196-
m[k] = v
197-
}
226+
func parseFlagsForConnectionOptions(m map[string]string, creds []*vault.Credential) []ConnectionOption {
198227
o := make([]ConnectionOption, 0)
199228
if apiEndpoint, ok := m["endpoint-url"]; ok {
200229
o = append(o, WithEndpoint(apiEndpoint))
@@ -300,20 +329,16 @@ func (p *AwsConnection) UpdateAsset(asset *inventory.Asset) {
300329
p.asset = asset
301330
}
302331

303-
func (p *AwsConnection) AccountId() string {
304-
return p.accountId
305-
}
306-
307332
func (p *AwsConnection) Profile() string {
308-
return p.profile
333+
return p.opts.profile
309334
}
310335

311336
func (p *AwsConnection) Scope() string {
312-
return p.scope
337+
return p.opts.scope
313338
}
314339

315340
func (p *AwsConnection) ConnectionOptions() map[string]string {
316-
return p.connectionOptions
341+
return p.opts.options
317342
}
318343

319344
func (p *AwsConnection) RunCommand(command string) (*shared.Command, error) {

‎providers/aws/connection/shared/shared.go‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
// Copyright (c) Mondoo, Inc.
2+
// SPDX-License-Identifier: BUSL-1.1
3+
4+
package shared
5+
6+
import (
7+
"go.mondoo.com/cnquery/v11/providers/os/connection/shared"
8+
)
9+
10+
type Connection interface {
11+
shared.Connection
12+
13+
// Used to avoid verifying a client with the same options more than once
14+
Verify() (accountID string, err error)
15+
Hash() uint64
16+
SetAccountId(string)
17+
}

‎providers/aws/go.mod‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -60,6 +60,7 @@ require (
6060
github.com/cockroachdb/errors v1.11.3
6161
github.com/google/uuid v1.6.0
6262
github.com/hashicorp/go-retryablehttp v0.7.7
63+
github.com/mitchellh/hashstructure/v2 v2.0.2
6364
github.com/rs/zerolog v1.34.0
6465
github.com/spf13/afero v1.14.0
6566
github.com/stretchr/testify v1.10.0
@@ -188,7 +189,7 @@ require (
188189
github.com/json-iterator/go v1.1.12 // indirect
189190
github.com/kballard/go-shellquote v0.0.0-20180428030007-95032a82bc51 // indirect
190191
github.com/kevinburke/ssh_config v1.2.0 // indirect
191-
github.com/klauspost/compress v1.17.11 // indirect
192+
github.com/klauspost/compress v1.18.0 // indirect
192193
github.com/knqyf263/go-rpmdb v0.1.1 // indirect
193194
github.com/kr/fs v0.1.0 // indirect
194195
github.com/kr/pretty v0.3.1 // indirect
@@ -198,7 +199,6 @@ require (
198199
github.com/mattn/go-colorable v0.1.14 // indirect
199200
github.com/mattn/go-isatty v0.0.20 // indirect
200201
github.com/mitchellh/go-homedir v1.1.0 // indirect
201-
github.com/mitchellh/hashstructure/v2 v2.0.2 // indirect
202202
github.com/mitchellh/mapstructure v1.5.0 // indirect
203203
github.com/moby/buildkit v0.19.0 // indirect
204204
github.com/moby/docker-image-spec v1.3.1 // indirect
@@ -221,9 +221,9 @@ require (
221221
github.com/pkg/sftp v1.13.9 // indirect
222222
github.com/planetscale/vtprotobuf v0.6.1-0.20240319094008-0393e58bdf10 // indirect
223223
github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
224-
github.com/prometheus/client_golang v1.20.5 // indirect
224+
github.com/prometheus/client_golang v1.22.0 // indirect
225225
github.com/prometheus/client_model v0.6.1 // indirect
226-
github.com/prometheus/common v0.55.0 // indirect
226+
github.com/prometheus/common v0.62.0 // indirect
227227
github.com/prometheus/procfs v0.15.1 // indirect
228228
github.com/remyoudompheng/bigfft v0.0.0-20230129092748-24d4a6f8daec // indirect
229229
github.com/rivo/uniseg v0.4.7 // indirect
@@ -278,13 +278,14 @@ require (
278278
gopkg.in/inf.v0 v0.9.1 // indirect
279279
gopkg.in/ini.v1 v1.67.0 // indirect
280280
gopkg.in/warnings.v0 v0.1.2 // indirect
281+
gopkg.in/yaml.v2 v2.4.0 // indirect
281282
gopkg.in/yaml.v3 v3.0.1 // indirect
282283
howett.net/plist v1.0.1 // indirect
283284
k8s.io/api v0.32.3 // indirect
284285
k8s.io/apimachinery v0.32.3 // indirect
285-
k8s.io/component-base v0.28.9 // indirect
286+
k8s.io/component-base v0.29.3 // indirect
286287
k8s.io/klog/v2 v2.130.1 // indirect
287-
k8s.io/kubelet v0.28.9 // indirect
288+
k8s.io/kubelet v0.29.3 // indirect
288289
k8s.io/utils v0.0.0-20250321185631-1f6e0b77f77e // indirect
289290
modernc.org/libc v1.62.1 // indirect
290291
modernc.org/mathutil v1.7.1 // indirect

‎providers/aws/go.sum‎

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -679,8 +679,8 @@ github.com/keybase/go-keychain v0.0.1 h1:way+bWYa6lDppZoZcgMbYsvC7GxljxrskdNInRt
679679
github.com/keybase/go-keychain v0.0.1/go.mod h1:PdEILRW3i9D8JcdM+FmY6RwkHGnhHxXwkPPMeUgOK1k=
680680
github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
681681
github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
682-
github.com/klauspost/compress v1.17.11 h1:In6xLpyWOi1+C7tXUUWv2ot1QvBjxevKAaI6IXrJmUc=
683-
github.com/klauspost/compress v1.17.11/go.mod h1:pMDklpSncoRMuLFrf1W9Ss9KT+0rH90U12bZKk7uwG0=
682+
github.com/klauspost/compress v1.18.0 h1:c/Cqfb0r+Yi+JtIEq73FWXVkRonBlf0CRNYc8Zttxdo=
683+
github.com/klauspost/compress v1.18.0/go.mod h1:2Pp+KzxcywXVXMr50+X0Q/Lsb43OQHYWRCY2AiWywWQ=
684684
github.com/knqyf263/go-rpmdb v0.1.1 h1:oh68mTCvp1XzxdU7EfafcWzzfstUZAEa3MW0IJye584=
685685
github.com/knqyf263/go-rpmdb v0.1.1/go.mod h1:9LQcoMCMQ9vrF7HcDtXfvqGO4+ddxFQ8+YF/0CVGDww=
686686
github.com/konsorten/go-windows-terminal-sequences v1.0.1/go.mod h1:T0+1ngSBFLxvqU3pZ+m/2kptfBszLMUkC4ZK/EgS/cQ=
@@ -815,8 +815,8 @@ github.com/posener/complete v1.2.3/go.mod h1:WZIdtGGp+qx0sLrYKtIRAruyNpv6hFCicSg
815815
github.com/prometheus/client_golang v0.9.1/go.mod h1:7SWBe2y4D6OKWSNQJUaRYU/AaXPKyh/dDVn+NZz0KFw=
816816
github.com/prometheus/client_golang v1.0.0/go.mod h1:db9x61etRT2tGnBNRi70OPL5FsnadC4Ky3P0J6CfImo=
817817
github.com/prometheus/client_golang v1.4.0/go.mod h1:e9GMxYsXl05ICDXkRhurwBS4Q3OK1iX/F2sw+iXX5zU=
818-
github.com/prometheus/client_golang v1.20.5 h1:cxppBPuYhUnsO6yo/aoRol4L7q7UFfdm+bR9r+8l63Y=
819-
github.com/prometheus/client_golang v1.20.5/go.mod h1:PIEt8X02hGcP8JWbeHyeZ53Y/jReSnHgO035n//V5WE=
818+
github.com/prometheus/client_golang v1.22.0 h1:rb93p9lokFEsctTys46VnV1kLCDpVZ0a/Y92Vm0Zc6Q=
819+
github.com/prometheus/client_golang v1.22.0/go.mod h1:R7ljNsLXhuQXYZYtw6GAE9AZg8Y7vEW5scdCXrWRXC0=
820820
github.com/prometheus/client_model v0.0.0-20180712105110-5c3871d89910/go.mod h1:MbSGuTsp3dbXC40dX6PRTWyKYBIrTGTE9sqQNg2J8bo=
821821
github.com/prometheus/client_model v0.0.0-20190129233127-fd36f4220a90/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
822822
github.com/prometheus/client_model v0.0.0-20190812154241-14fe0d1b01d4/go.mod h1:xMI15A0UPsDsEKsMN9yxemIoYk6Tm2C1GtYGdfGttqA=
@@ -825,8 +825,8 @@ github.com/prometheus/client_model v0.6.1 h1:ZKSh/rekM+n3CeS952MLRAdFwIKqeY8b62p
825825
github.com/prometheus/client_model v0.6.1/go.mod h1:OrxVMOVHjw3lKMa8+x6HeMGkHMQyHDk9E3jmP2AmGiY=
826826
github.com/prometheus/common v0.4.1/go.mod h1:TNfzLD0ON7rHzMJeJkieUDPYmFC7Snx/y86RQel1bk4=
827827
github.com/prometheus/common v0.9.1/go.mod h1:yhUN8i9wzaXS3w1O07YhxHEBxD+W35wd8bs7vj7HSQ4=
828-
github.com/prometheus/common v0.55.0 h1:KEi6DK7lXW/m7Ig5i47x0vRzuBsHuvJdi5ee6Y3G1dc=
829-
github.com/prometheus/common v0.55.0/go.mod h1:2SECS4xJG1kd8XF9IcM1gMX6510RAEL65zxzNImwdc8=
828+
github.com/prometheus/common v0.62.0 h1:xasJaQlnWAeyHdUBeGjXmutelfJHWMRr+Fg4QszZ2Io=
829+
github.com/prometheus/common v0.62.0/go.mod h1:vyBcEuLSvWos9B1+CyL7JZ2up+uFzXhkqml0W5zIY1I=
830830
github.com/prometheus/procfs v0.0.0-20181005140218-185b4288413d/go.mod h1:c3At6R/oaqEKCNdg8wHV1ftS6bRYblBhIjjI8uT2IGk=
831831
github.com/prometheus/procfs v0.0.2/go.mod h1:TjEm7ze935MbeOT/UhFTIMYKhuLP4wbCsTZCD3I8kEA=
832832
github.com/prometheus/procfs v0.0.8/go.mod h1:7Qr8sr6344vo1JqZ6HhLceV9o3AJ1Ff+GxbHq6oeK9A=
@@ -1511,6 +1511,7 @@ gopkg.in/yaml.v2 v2.2.3/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
15111511
gopkg.in/yaml.v2 v2.2.4/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
15121512
gopkg.in/yaml.v2 v2.2.5/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
15131513
gopkg.in/yaml.v2 v2.2.8/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
1514+
gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
15141515
gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
15151516
gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
15161517
gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
@@ -1533,12 +1534,12 @@ k8s.io/apimachinery v0.32.3 h1:JmDuDarhDmA/Li7j3aPrwhpNBA94Nvk5zLeOge9HH1U=
15331534
k8s.io/apimachinery v0.32.3/go.mod h1:GpHVgxoKlTxClKcteaeuF1Ul/lDVb74KpZcxcmLDElE=
15341535
k8s.io/client-go v0.32.3 h1:RKPVltzopkSgHS7aS98QdscAgtgah/+zmpAogooIqVU=
15351536
k8s.io/client-go v0.32.3/go.mod h1:3v0+3k4IcT9bXTc4V2rt+d2ZPPG700Xy6Oi0Gdl2PaY=
1536-
k8s.io/component-base v0.28.9 h1:ySM2PR8Z/xaUSG1Akd3yM6dqUezTltI7S5aV41MMuuc=
1537-
k8s.io/component-base v0.28.9/go.mod h1:QtWzscEhCKRfHV24/S+11BwWjVxhC6fd3RYoEgZcWFU=
1537+
k8s.io/component-base v0.29.3 h1:Oq9/nddUxlnrCuuR2K/jp6aflVvc0uDvxMzAWxnGzAo=
1538+
k8s.io/component-base v0.29.3/go.mod h1:Yuj33XXjuOk2BAaHsIGHhCKZQAgYKhqIxIjIr2UXYio=
15381539
k8s.io/klog/v2 v2.130.1 h1:n9Xl7H1Xvksem4KFG4PYbdQCQxqc/tTUyrgXaOhHSzk=
15391540
k8s.io/klog/v2 v2.130.1/go.mod h1:3Jpz1GvMt720eyJH1ckRHK1EDfpxISzJ7I9OYgaDtPE=
1540-
k8s.io/kubelet v0.28.9 h1:76v00fFLeniz27kXhGGUIxONdwa9LKcD2Jd5cXYAZko=
1541-
k8s.io/kubelet v0.28.9/go.mod h1:46P39DFjI+E59nU2OgpatyS3oWy58ClulKO6riZ/97o=
1541+
k8s.io/kubelet v0.29.3 h1:X9h0ZHzc+eUeNTaksbN0ItHyvGhQ7Z0HPjnQD2oHdwU=
1542+
k8s.io/kubelet v0.29.3/go.mod h1:jDiGuTkFOUynyBKzOoC1xRSWlgAZ9UPcTYeFyjr6vas=
15421543
k8s.io/utils v0.0.0-20250321185631-1f6e0b77f77e h1:KqK5c/ghOm8xkHYhlodbp6i6+r+ChV2vuAuVRdFbLro=
15431544
k8s.io/utils v0.0.0-20250321185631-1f6e0b77f77e/go.mod h1:OLgZIPagt7ERELqWJFomSt595RzquPNLL48iOWgYOg0=
15441545
modernc.org/cc/v4 v4.25.2 h1:T2oH7sZdGvTaie0BRNFbIYsabzCxUQg8nLqCdQ2i0ic=

0 commit comments

Comments
 (0)