Skip to content

Commit d673bdb

Browse files
agneumagneum
committed
Merge branch 'timing-estimator' into 'master'
feat: add timing estimator (#212) # Description - move the estimator package to the DBLab repository in order to have an ability to use it in CI observer as well - provide web-socket API to start estimation remotely (for example, from the Joe Bot repository) # Related issue #212 # Examples https://gitlab.com/postgres-ai/database-lab/-/issues/212#note_520279684 See merge request postgres-ai/database-lab!270
2 parents cc8a8f7 + a0e7342 commit d673bdb

21 files changed

+1117
-8
lines changed

‎cmd/database-lab/main.go

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ import (
2121
"github.com/rs/xid"
2222

2323
"gitlab.com/postgres-ai/database-lab/v2/pkg/config"
24+
"gitlab.com/postgres-ai/database-lab/v2/pkg/estimator"
2425
"gitlab.com/postgres-ai/database-lab/v2/pkg/log"
2526
"gitlab.com/postgres-ai/database-lab/v2/pkg/observer"
2627
"gitlab.com/postgres-ai/database-lab/v2/pkg/retrieval"
@@ -108,18 +109,19 @@ func main() {
108109
}
109110

110111
obs := observer.NewObserver(dockerCLI, &cfg.Observer, platformSvc.Client, pm)
112+
est := estimator.NewEstimator(&cfg.Estimator)
111113

112114
go removeObservingClones(obsCh, obs)
113115

114-
server := srv.NewServer(&cfg.Server, obs, cloningSvc, platformSvc, dockerCLI)
116+
server := srv.NewServer(&cfg.Server, obs, cloningSvc, platformSvc, dockerCLI, est)
115117

116118
reloadCh := setReloadListener()
117119

118120
go func() {
119121
for range reloadCh {
120122
log.Msg("Reloading configuration")
121123

122-
if err := reloadConfig(ctx, instanceID, provisionSvc, retrievalSvc, pm, cloningSvc, platformSvc, server); err != nil {
124+
if err := reloadConfig(ctx, instanceID, provisionSvc, retrievalSvc, pm, cloningSvc, platformSvc, est, server); err != nil {
123125
log.Err("Failed to reload configuration", err)
124126
}
125127

@@ -166,7 +168,7 @@ func loadConfiguration(instanceID string) (*config.Config, error) {
166168
}
167169

168170
func reloadConfig(ctx context.Context, instanceID string, provisionSvc *provision.Provisioner, retrievalSvc *retrieval.Retrieval,
169-
pm *pool.Manager, cloningSvc cloning.Cloning, platformSvc *platform.Service, server *srv.Server) error {
171+
pm *pool.Manager, cloningSvc cloning.Cloning, platformSvc *platform.Service, est *estimator.Estimator, server *srv.Server) error {
170172
cfg, err := loadConfiguration(instanceID)
171173
if err != nil {
172174
return err
@@ -198,6 +200,7 @@ func reloadConfig(ctx context.Context, instanceID string, provisionSvc *provisio
198200
retrievalSvc.Reload(ctx, cfg)
199201
cloningSvc.Reload(cfg.Cloning)
200202
platformSvc.Reload(newPlatformSvc)
203+
est.Reload(cfg.Estimator)
201204
server.Reload(cfg.Server)
202205

203206
return nil

‎configs/config.example.logical_generic.yml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,3 +296,17 @@ cloning:
296296
# "regexp": "replace"
297297
# "select \\d+": "***"
298298
# "[a-z0-9._%+\\-]+(@[a-z0-9.\\-]+\\.[a-z]{2,4})": "***$1"
299+
#
300+
# Tool to calculate timing difference between Database Lab and production environments.
301+
#estimator:
302+
# # The ratio evaluating the timing difference for operations involving IO Read between Database Lab and production environments.
303+
# readRatio: 1
304+
#
305+
# # The ratio evaluating the timing difference for operations involving IO Write between Database Lab and production environments.
306+
# writeRatio: 1
307+
#
308+
# # Time interval of samples taken by the profiler.
309+
# profilingInterval: 10ms
310+
#
311+
# # The minimum number of samples sufficient to display the estimation results.
312+
# sampleThreshold: 20

‎configs/config.example.logical_rds_iam.yml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -296,3 +296,17 @@ cloning:
296296
# "regexp": "replace"
297297
# "select \\d+": "***"
298298
# "[a-z0-9._%+\\-]+(@[a-z0-9.\\-]+\\.[a-z]{2,4})": "***$1"
299+
#
300+
# Tool to calculate timing difference between Database Lab and production environments.
301+
#estimator:
302+
# # The ratio evaluating the timing difference for operations involving IO Read between Database Lab and production environments.
303+
# readRatio: 1
304+
#
305+
# # The ratio evaluating the timing difference for operations involving IO Write between Database Lab and production environments.
306+
# writeRatio: 1
307+
#
308+
# # Time interval of samples taken by the profiler.
309+
# profilingInterval: 10ms
310+
#
311+
# # The minimum number of samples sufficient to display the estimation results.
312+
# sampleThreshold: 20

‎configs/config.example.physical_generic.yml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -313,3 +313,17 @@ cloning:
313313
# "regexp": "replace"
314314
# "select \\d+": "***"
315315
# "[a-z0-9._%+\\-]+(@[a-z0-9.\\-]+\\.[a-z]{2,4})": "***$1"
316+
#
317+
# Tool to calculate timing difference between Database Lab and production environments.
318+
#estimator:
319+
# # The ratio evaluating the timing difference for operations involving IO Read between Database Lab and production environments.
320+
# readRatio: 1
321+
#
322+
# # The ratio evaluating the timing difference for operations involving IO Write between Database Lab and production environments.
323+
# writeRatio: 1
324+
#
325+
# # Time interval of samples taken by the profiler.
326+
# profilingInterval: 10ms
327+
#
328+
# # The minimum number of samples sufficient to display the estimation results.
329+
# sampleThreshold: 20

‎configs/config.example.physical_walg.yml

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -301,3 +301,17 @@ cloning:
301301
# "regexp": "replace"
302302
# "select \\d+": "***"
303303
# "[a-z0-9._%+\\-]+(@[a-z0-9.\\-]+\\.[a-z]{2,4})": "***$1"
304+
#
305+
# Tool to calculate timing difference between Database Lab and production environments.
306+
#estimator:
307+
# # The ratio evaluating the timing difference for operations involving IO Read between Database Lab and production environments.
308+
# readRatio: 1
309+
#
310+
# # The ratio evaluating the timing difference for operations involving IO Write between Database Lab and production environments.
311+
# writeRatio: 1
312+
#
313+
# # Time interval of samples taken by the profiler.
314+
# profilingInterval: 10ms
315+
#
316+
# # The minimum number of samples sufficient to display the estimation results.
317+
# sampleThreshold: 20

‎go.mod

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,8 @@ require (
1919
github.com/go-ole/go-ole v1.2.4 // indirect
2020
github.com/gogo/protobuf v1.3.1 // indirect
2121
github.com/gorilla/mux v1.8.0
22+
github.com/gorilla/websocket v1.4.2
23+
github.com/jackc/pgtype v1.5.0
2224
github.com/jackc/pgx/v4 v4.9.0
2325
github.com/lib/pq v1.8.0
2426
github.com/morikuni/aec v1.0.0 // indirect

‎go.sum

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,8 @@ github.com/google/go-cmp v0.4.0/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/
6464
github.com/google/renameio v0.1.0/go.mod h1:KWCgfxg9yswjAJkECMjeO8J8rahYeXnNhOm40UhjYkI=
6565
github.com/gorilla/mux v1.8.0 h1:i40aqfkR1h2SlN9hojwV5ZA91wcXFOvkdNIeFDP5koI=
6666
github.com/gorilla/mux v1.8.0/go.mod h1:DVbg23sWSpFRCP0SfiEN6jmj59UnW/n46BH5rLB71So=
67+
github.com/gorilla/websocket v1.4.2 h1:+/TMaTYc4QFitKJxsQ7Yye35DkWvkdLcvGKqM+x0Ufc=
68+
github.com/gorilla/websocket v1.4.2/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
6769
github.com/jackc/chunkreader v1.0.0 h1:4s39bBR8ByfqH+DKm8rQA3E1LHZWB9XWcrz8fqaZbe0=
6870
github.com/jackc/chunkreader v1.0.0/go.mod h1:RT6O25fNZIuasFJRyZ4R/Y2BbhasbmZXF9QQ7T3kePo=
6971
github.com/jackc/chunkreader/v2 v2.0.0/go.mod h1:odVSm741yZoC3dpHEUXIqA9tQRhFrgOHwnPIn9lDKlk=

‎pkg/client/dblabapi/estimator.go

Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
/*
2+
2021 © Postgres.ai
3+
*/
4+
5+
package dblabapi
6+
7+
import (
8+
"context"
9+
"encoding/json"
10+
"net/url"
11+
12+
"github.com/gorilla/websocket"
13+
"github.com/pkg/errors"
14+
15+
"gitlab.com/postgres-ai/database-lab/v2/pkg/estimator"
16+
"gitlab.com/postgres-ai/database-lab/v2/pkg/log"
17+
)
18+
19+
// EstimatorClient defines a client to perform timing estimation.
20+
type EstimatorClient struct {
21+
conn *websocket.Conn
22+
results chan estimator.Result
23+
ready chan struct{}
24+
}
25+
26+
// Close closes connection.
27+
func (e *EstimatorClient) Close() error {
28+
return e.conn.Close()
29+
}
30+
31+
// Wait waits for connection readiness.
32+
func (e *EstimatorClient) Wait() chan struct{} {
33+
return e.ready
34+
}
35+
36+
// ReadResult returns estimation results.
37+
func (e *EstimatorClient) ReadResult() estimator.Result {
38+
return <-e.results
39+
}
40+
41+
// SetReadBlocks sends a number of read blocks.
42+
func (e *EstimatorClient) SetReadBlocks(readBlocks uint64) error {
43+
result := estimator.ReadBlocksEvent{
44+
EventType: estimator.ReadBlocksType,
45+
ReadBlocks: readBlocks,
46+
}
47+
48+
readBlocksData, err := json.Marshal(result)
49+
if err != nil {
50+
return err
51+
}
52+
53+
return e.conn.WriteMessage(websocket.TextMessage, readBlocksData)
54+
}
55+
56+
// Estimate creates connection for estimation session.
57+
func (c *Client) Estimate(ctx context.Context, cloneID, pid string) (*EstimatorClient, error) {
58+
u := c.URL("/estimate")
59+
u.Scheme = "ws"
60+
61+
values := url.Values{}
62+
values.Add("clone_id", cloneID)
63+
values.Add("pid", pid)
64+
u.RawQuery = values.Encode()
65+
66+
log.Dbg("connecting to ", u.String())
67+
68+
conn, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
69+
if err != nil {
70+
return nil, errors.Wrap(err, "failed to connect")
71+
}
72+
73+
est := &EstimatorClient{
74+
conn: conn,
75+
results: make(chan estimator.Result, 1),
76+
ready: make(chan struct{}, 1),
77+
}
78+
79+
go func() {
80+
if err := est.receiveMessages(ctx); err != nil {
81+
log.Dbg("error while receive messages: ", err)
82+
}
83+
}()
84+
85+
return est, nil
86+
}
87+
88+
func (e *EstimatorClient) receiveMessages(ctx context.Context) error {
89+
for {
90+
if ctx.Err() != nil {
91+
log.Msg(ctx.Err())
92+
break
93+
}
94+
95+
_, message, err := e.conn.ReadMessage()
96+
if err != nil {
97+
return err
98+
}
99+
100+
event := estimator.Event{}
101+
if err := json.Unmarshal(message, &event); err != nil {
102+
return err
103+
}
104+
105+
switch event.EventType {
106+
case estimator.ReadyEventType:
107+
e.ready <- struct{}{}
108+
109+
case estimator.ResultEventType:
110+
result := estimator.ResultEvent{}
111+
if err := json.Unmarshal(message, &result); err != nil {
112+
log.Dbg("failed to read the result event: ", err)
113+
break
114+
}
115+
116+
e.results <- result.Payload
117+
if err := e.conn.Close(); err != nil {
118+
log.Dbg("failed to close connection: ", err)
119+
}
120+
121+
return nil
122+
}
123+
124+
log.Dbg("received unknown event type: ", event.EventType, string(message))
125+
}
126+
127+
return nil
128+
}

‎pkg/config/config.go

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ import (
1111
"github.com/pkg/errors"
1212
"gopkg.in/yaml.v2"
1313

14+
"gitlab.com/postgres-ai/database-lab/v2/pkg/estimator"
1415
"gitlab.com/postgres-ai/database-lab/v2/pkg/observer"
1516
retConfig "gitlab.com/postgres-ai/database-lab/v2/pkg/retrieval/config"
1617
"gitlab.com/postgres-ai/database-lab/v2/pkg/retrieval/engine/postgres/tools/defaults"
@@ -31,6 +32,7 @@ type Config struct {
3132
Global Global `yaml:"global"`
3233
Retrieval retConfig.Config `yaml:"retrieval"`
3334
Observer observer.Config `yaml:"observer"`
35+
Estimator estimator.Config `yaml:"estimator"`
3436
PoolManager pool.Config `yaml:"poolManager"`
3537
}
3638

‎pkg/estimator/estimator.go

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
/*
2+
2021 © Postgres.ai
3+
*/
4+
5+
// Package estimator provides tools to estimate query timing for a production environment.
6+
package estimator
7+
8+
import (
9+
"context"
10+
"time"
11+
)
12+
13+
// Config describes options to estimate query timing.
14+
type Config struct {
15+
ReadRatio float64 `yaml:"readRatio"`
16+
WriteRatio float64 `yaml:"writeRatio"`
17+
ProfilingInterval time.Duration `yaml:"profilingInterval"`
18+
SampleThreshold int `yaml:"sampleThreshold"`
19+
}
20+
21+
// Estimator defines a timing estimator.
22+
type Estimator struct {
23+
cfg *Config
24+
}
25+
26+
// NewEstimator creates a new Estimator.
27+
func NewEstimator(cfg *Config) *Estimator {
28+
return &Estimator{cfg: cfg}
29+
}
30+
31+
// Run starts profiling if it needs to be done.
32+
func (e *Estimator) Run(ctx context.Context, p *Profiler) {
33+
if p.opts.SampleThreshold > 0 && shouldEstimate(e.cfg.ReadRatio, e.cfg.WriteRatio) {
34+
go p.Start(ctx)
35+
return
36+
}
37+
38+
p.Stop()
39+
}
40+
41+
// Reload reloads estimator configuration.
42+
func (e *Estimator) Reload(cfg Config) {
43+
*e.cfg = cfg
44+
}
45+
46+
// Config returns Estimator configuration.
47+
func (e *Estimator) Config() Config {
48+
if e.cfg == nil {
49+
return Config{}
50+
}
51+
52+
return *e.cfg
53+
}
54+
55+
// shouldEstimate checks ratios to determine whether to skip an estimation.
56+
func shouldEstimate(readRatio, writeRatio float64) bool {
57+
return (readRatio != 0 || writeRatio != 0) && (readRatio != 1 || writeRatio != 1)
58+
}

0 commit comments

Comments
 (0)