forked from vitessio/vitess
-
Notifications
You must be signed in to change notification settings - Fork 13
/
dbconn.go
213 lines (192 loc) · 6.37 KB
/
dbconn.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package tabletserver
import (
"errors"
"fmt"
"time"
log "github.com/golang/glog"
mproto "github.com/youtube/vitess/go/mysql/proto"
"github.com/youtube/vitess/go/sqldb"
"github.com/youtube/vitess/go/sync2"
"github.com/youtube/vitess/go/vt/dbconnpool"
"github.com/youtube/vitess/go/vt/proto/vtrpc"
"golang.org/x/net/context"
)
// DBConn is a db connection for tabletserver.
// It performs automatic reconnects as needed.
// Its Execute function has a timeout that can kill
// its own queries and the underlying connection.
// It will also trigger a CheckMySQL whenever applicable.
type DBConn struct {
conn *dbconnpool.DBConnection
info *sqldb.ConnParams
pool *ConnPool
queryServiceStats *QueryServiceStats
current sync2.AtomicString
}
// NewDBConn creates a new DBConn. It triggers a CheckMySQL if creation fails.
func NewDBConn(
cp *ConnPool,
appParams,
dbaParams *sqldb.ConnParams,
qStats *QueryServiceStats) (*DBConn, error) {
c, err := dbconnpool.NewDBConnection(appParams, qStats.MySQLStats)
if err != nil {
cp.checker.CheckMySQL()
return nil, err
}
return &DBConn{
conn: c,
info: appParams,
pool: cp,
queryServiceStats: qStats,
}, nil
}
// Exec executes the specified query. If there is a connection error, it will reconnect
// and retry. A failed reconnect will trigger a CheckMySQL.
func (dbc *DBConn) Exec(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
for attempt := 1; attempt <= 2; attempt++ {
r, err := dbc.execOnce(ctx, query, maxrows, wantfields)
switch {
case err == nil:
return r, nil
case !IsConnErr(err):
// MySQL error that isn't due to a connection issue
return nil, NewTabletErrorSQL(ErrFail, vtrpc.ErrorCode_UNKNOWN_ERROR, err)
case attempt == 2:
// If the MySQL connection is bad, we assume that there is nothing wrong with
// the query itself, and retrying it might succeed. The MySQL connection might
// fix itself, or the query could succeed on a different VtTablet.
return nil, NewTabletErrorSQL(ErrFatal, vtrpc.ErrorCode_INTERNAL_ERROR, err)
}
err2 := dbc.reconnect()
if err2 != nil {
dbc.pool.checker.CheckMySQL()
return nil, NewTabletErrorSQL(ErrFatal, vtrpc.ErrorCode_INTERNAL_ERROR, err)
}
}
return nil, NewTabletErrorSQL(ErrFatal, vtrpc.ErrorCode_INTERNAL_ERROR, errors.New("dbconn.Exec: unreachable code"))
}
func (dbc *DBConn) execOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
dbc.current.Set(query)
defer dbc.current.Set("")
done := dbc.setDeadline(ctx)
if done != nil {
defer close(done)
}
// Uncomment this line for manual testing.
// defer time.Sleep(20 * time.Second)
return dbc.conn.ExecuteFetch(query, maxrows, wantfields)
}
// ExecOnce executes the specified query, but does not retry on connection errors.
func (dbc *DBConn) ExecOnce(ctx context.Context, query string, maxrows int, wantfields bool) (*mproto.QueryResult, error) {
return dbc.execOnce(ctx, query, maxrows, wantfields)
}
// Stream executes the query and streams the results.
func (dbc *DBConn) Stream(ctx context.Context, query string, callback func(*mproto.QueryResult) error, streamBufferSize int) error {
dbc.current.Set(query)
defer dbc.current.Set("")
done := dbc.setDeadline(ctx)
if done != nil {
defer close(done)
}
return dbc.conn.ExecuteStreamFetch(query, callback, streamBufferSize)
}
// VerifyStrict returns true if MySQL is in STRICT mode.
func (dbc *DBConn) VerifyStrict() bool {
return dbc.conn.VerifyStrict()
}
// Close closes the DBConn.
func (dbc *DBConn) Close() {
dbc.conn.Close()
}
// IsClosed returns true if DBConn is closed.
func (dbc *DBConn) IsClosed() bool {
return dbc.conn.IsClosed()
}
// Recycle returns the DBConn to the pool.
func (dbc *DBConn) Recycle() {
if dbc.conn.IsClosed() {
dbc.pool.Put(nil)
} else {
dbc.pool.Put(dbc)
}
}
// Kill kills the currently executing query both on MySQL side
// and on the connection side. If no query is executing, it's a no-op.
// Kill will also not kill a query more than once.
func (dbc *DBConn) Kill() error {
dbc.queryServiceStats.KillStats.Add("Queries", 1)
log.Infof("killing query %s", dbc.Current())
killConn, err := dbc.pool.dbaPool.Get(0)
if err != nil {
log.Warningf("Failed to get conn from dba pool: %v", err)
// TODO(aaijazi): Find the right error code for an internal error that we don't want to retry
return NewTabletError(ErrFail, vtrpc.ErrorCode_INTERNAL_ERROR, "Failed to get conn from dba pool: %v", err)
}
defer killConn.Recycle()
sql := fmt.Sprintf("kill %d", dbc.conn.ID())
_, err = killConn.ExecuteFetch(sql, 10000, false)
if err != nil {
log.Errorf("Could not kill query %s: %v", dbc.Current(), err)
// TODO(aaijazi): Find the right error code for an internal error that we don't want to retry
return NewTabletError(ErrFail, vtrpc.ErrorCode_INTERNAL_ERROR, "Could not kill query %s: %v", dbc.Current(), err)
}
return nil
}
// Current returns the currently executing query.
func (dbc *DBConn) Current() string {
return dbc.current.Get()
}
// ID returns the connection id.
func (dbc *DBConn) ID() int64 {
return dbc.conn.ID()
}
func (dbc *DBConn) reconnect() error {
dbc.conn.Close()
newConn, err := dbconnpool.NewDBConnection(dbc.info, dbc.queryServiceStats.MySQLStats)
if err != nil {
return err
}
dbc.conn = newConn
return nil
}
func (dbc *DBConn) setDeadline(ctx context.Context) chan bool {
if ctx.Done() == nil {
return nil
}
done := make(chan bool)
go func() {
startTime := time.Now()
select {
case <-ctx.Done():
// There is a possibility that the query returned very fast,
// which will cause ctx to get canceled. Check for this condition.
select {
case <-done:
return
default:
}
dbc.Kill()
case <-done:
return
}
elapsed := time.Now().Sub(startTime)
// Give 2x the elapsed time and some buffer as grace period
// for the query to get killed.
tmr2 := time.NewTimer(2*elapsed + 5*time.Second)
defer tmr2.Stop()
select {
case <-tmr2.C:
dbc.queryServiceStats.InternalErrors.Add("HungQuery", 1)
log.Warningf("Query may be hung: %s", dbc.Current())
case <-done:
return
}
<-done
log.Warningf("Hung query returned")
}()
return done
}