package tidb import ( "context" "database/sql" "fmt" "sync" "sync/atomic" "time" "github.com/bilibili/kratos/pkg/log" "github.com/bilibili/kratos/pkg/naming" "github.com/bilibili/kratos/pkg/net/netutil/breaker" "github.com/bilibili/kratos/pkg/net/trace" "github.com/go-sql-driver/mysql" "github.com/pkg/errors" ) const ( _family = "tidb_client" _slowLogDuration = time.Millisecond * 250 ) var ( // ErrStmtNil prepared stmt error ErrStmtNil = errors.New("sql: prepare failed and stmt nil") // ErrNoRows is returned by Scan when QueryRow doesn't return a row. // In such a case, QueryRow returns a placeholder *Row value that defers // this error until a Scan. ErrNoRows = sql.ErrNoRows // ErrTxDone transaction done. ErrTxDone = sql.ErrTxDone ) // DB database. type DB struct { conf *Config conns []*conn idx int64 dis naming.Resolver appid string mutex sync.RWMutex breakerGroup *breaker.Group } // conn database connection type conn struct { *sql.DB breaker breaker.Breaker conf *Config addr string } // Tx transaction. type Tx struct { db *conn tx *sql.Tx t trace.Trace c context.Context cancel func() } // Row row. type Row struct { err error *sql.Row db *conn query string args []interface{} t trace.Trace cancel func() } // Scan copies the columns from the matched row into the values pointed at by dest. func (r *Row) Scan(dest ...interface{}) (err error) { defer slowLog(fmt.Sprintf("Scan addr: %s query(%s) args(%+v)", r.db.addr, r.query, r.args), time.Now()) if r.t != nil { defer r.t.Finish(&err) } if r.err != nil { err = r.err } else if r.Row == nil { err = ErrStmtNil } if err != nil { return } err = r.Row.Scan(dest...) if r.cancel != nil { r.cancel() } r.db.onBreaker(&err) if err != ErrNoRows { err = errors.Wrapf(err, "addr: %s, query %s args %+v", r.db.addr, r.query, r.args) } return } // Rows rows. type Rows struct { *sql.Rows cancel func() } // Close closes the Rows, preventing further enumeration. If Next is called // and returns false and there are no further result sets, // the Rows are closed automatically and it will suffice to check the // result of Err. Close is idempotent and does not affect the result of Err. func (rs *Rows) Close() (err error) { err = errors.WithStack(rs.Rows.Close()) if rs.cancel != nil { rs.cancel() } return } // Stmt prepared stmt. type Stmt struct { db *conn tx bool query string stmt atomic.Value t trace.Trace } // Stmts random prepared stmt. type Stmts struct { query string sts map[string]*Stmt mu sync.RWMutex db *DB } // Open opens a database specified by its database driver name and a // driver-specific data source name, usually consisting of at least a database // name and connection information. func Open(c *Config) (db *DB, err error) { db = &DB{conf: c, breakerGroup: breaker.NewGroup(c.Breaker)} cfg, err := mysql.ParseDSN(c.DSN) if err != nil { return } var dsns []string if cfg.Net == "discovery" { db.appid = cfg.Addr for _, addr := range db.disc() { dsns = append(dsns, genDSN(c.DSN, addr)) } } else { dsns = append(dsns, c.DSN) } cs := make([]*conn, 0, len(dsns)) for _, dsn := range dsns { r, err := db.connectDSN(dsn) if err != nil { return db, err } cs = append(cs, r) } db.conns = cs return } func (db *DB) connectDSN(dsn string) (c *conn, err error) { d, err := connect(db.conf, dsn) if err != nil { return } addr := parseDSNAddr(dsn) brk := db.breakerGroup.Get(addr) c = &conn{DB: d, breaker: brk, conf: db.conf, addr: addr} return } func connect(c *Config, dataSourceName string) (*sql.DB, error) { d, err := sql.Open("mysql", dataSourceName) if err != nil { err = errors.WithStack(err) return nil, err } d.SetMaxOpenConns(c.Active) d.SetMaxIdleConns(c.Idle) d.SetConnMaxLifetime(time.Duration(c.IdleTimeout)) return d, nil } func (db *DB) conn() (c *conn) { db.mutex.RLock() c = db.conns[db.index()] db.mutex.RUnlock() return } // Begin starts a transaction. The isolation level is dependent on the driver. func (db *DB) Begin(c context.Context) (tx *Tx, err error) { return db.conn().begin(c) } // Exec executes a query without returning any rows. // The args are for any placeholder parameters in the query. func (db *DB) Exec(c context.Context, query string, args ...interface{}) (res sql.Result, err error) { return db.conn().exec(c, query, args...) } // Prepare creates a prepared statement for later queries or executions. // Multiple queries or executions may be run concurrently from the returned // statement. The caller must call the statement's Close method when the // statement is no longer needed. func (db *DB) Prepare(query string) (*Stmt, error) { return db.conn().prepare(query) } // Prepared creates a prepared statement for later queries or executions. // Multiple queries or executions may be run concurrently from the returned // statement. The caller must call the statement's Close method when the // statement is no longer needed. func (db *DB) Prepared(query string) (s *Stmts) { s = &Stmts{query: query, sts: make(map[string]*Stmt), db: db} for _, c := range db.conns { st := c.prepared(query) s.mu.Lock() s.sts[c.addr] = st s.mu.Unlock() } return } // Query executes a query that returns rows, typically a SELECT. The args are // for any placeholder parameters in the query. func (db *DB) Query(c context.Context, query string, args ...interface{}) (rows *Rows, err error) { return db.conn().query(c, query, args...) } // QueryRow executes a query that is expected to return at most one row. // QueryRow always returns a non-nil value. Errors are deferred until Row's // Scan method is called. func (db *DB) QueryRow(c context.Context, query string, args ...interface{}) *Row { return db.conn().queryRow(c, query, args...) } func (db *DB) index() int { if len(db.conns) == 1 { return 0 } v := atomic.AddInt64(&db.idx, 1) return int(v) % len(db.conns) } // Close closes the databases, releasing any open resources. func (db *DB) Close() (err error) { db.mutex.RLock() defer db.mutex.RUnlock() for _, d := range db.conns { if e := d.Close(); e != nil { err = errors.WithStack(e) } } return } // Ping verifies a connection to the database is still alive, establishing a // connection if necessary. func (db *DB) Ping(c context.Context) (err error) { if err = db.conn().ping(c); err != nil { return } return } func (db *conn) onBreaker(err *error) { if err != nil && *err != nil && *err != sql.ErrNoRows && *err != sql.ErrTxDone { db.breaker.MarkFailed() } else { db.breaker.MarkSuccess() } } func (db *conn) begin(c context.Context) (tx *Tx, err error) { now := time.Now() defer slowLog(fmt.Sprintf("Begin addr: %s", db.addr), now) t, ok := trace.FromContext(c) if ok { t = t.Fork(_family, "begin") t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, "")) defer func() { if err != nil { t.Finish(&err) } }() } if err = db.breaker.Allow(); err != nil { _metricReqErr.Inc(db.addr, db.addr, "begin", "breaker") return } _, c, cancel := db.conf.TranTimeout.Shrink(c) rtx, err := db.BeginTx(c, nil) _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), db.addr, db.addr, "begin") if err != nil { err = errors.WithStack(err) cancel() return } tx = &Tx{tx: rtx, t: t, db: db, c: c, cancel: cancel} return } func (db *conn) exec(c context.Context, query string, args ...interface{}) (res sql.Result, err error) { now := time.Now() defer slowLog(fmt.Sprintf("Exec addr: %s query(%s) args(%+v)", db.addr, query, args), now) if t, ok := trace.FromContext(c); ok { t = t.Fork(_family, "exec") t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, query)) defer t.Finish(&err) } if err = db.breaker.Allow(); err != nil { _metricReqErr.Inc(db.addr, db.addr, "exec", "breaker") return } _, c, cancel := db.conf.ExecTimeout.Shrink(c) res, err = db.ExecContext(c, query, args...) cancel() db.onBreaker(&err) _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), db.addr, db.addr, "exec") if err != nil { err = errors.Wrapf(err, "addr: %s exec:%s, args:%+v", db.addr, query, args) } return } func (db *conn) ping(c context.Context) (err error) { now := time.Now() defer slowLog(fmt.Sprintf("Ping addr: %s", db.addr), now) if t, ok := trace.FromContext(c); ok { t = t.Fork(_family, "ping") t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, "")) defer t.Finish(&err) } if err = db.breaker.Allow(); err != nil { _metricReqErr.Inc(db.addr, db.addr, "ping", "breaker") return } _, c, cancel := db.conf.ExecTimeout.Shrink(c) err = db.PingContext(c) cancel() db.onBreaker(&err) _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), db.addr, db.addr, "ping") if err != nil { err = errors.WithStack(err) } return } func (db *conn) prepare(query string) (*Stmt, error) { defer slowLog(fmt.Sprintf("Prepare addr: %s query(%s)", db.addr, query), time.Now()) stmt, err := db.Prepare(query) if err != nil { err = errors.Wrapf(err, "addr: %s prepare %s", db.addr, query) return nil, err } st := &Stmt{query: query, db: db} st.stmt.Store(stmt) return st, nil } func (db *conn) prepared(query string) (stmt *Stmt) { defer slowLog(fmt.Sprintf("Prepared addr: %s query(%s)", db.addr, query), time.Now()) stmt = &Stmt{query: query, db: db} s, err := db.Prepare(query) if err == nil { stmt.stmt.Store(s) return } return } func (db *conn) query(c context.Context, query string, args ...interface{}) (rows *Rows, err error) { now := time.Now() defer slowLog(fmt.Sprintf("Query addr: %s query(%s) args(%+v)", db.addr, query, args), now) if t, ok := trace.FromContext(c); ok { t = t.Fork(_family, "query") t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, query)) defer t.Finish(&err) } if err = db.breaker.Allow(); err != nil { _metricReqErr.Inc(db.addr, db.addr, "query", "breaker") return } _, c, cancel := db.conf.QueryTimeout.Shrink(c) rs, err := db.DB.QueryContext(c, query, args...) db.onBreaker(&err) _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), db.addr, db.addr, "query") if err != nil { err = errors.Wrapf(err, "addr: %s, query:%s, args:%+v", db.addr, query, args) cancel() return } rows = &Rows{Rows: rs, cancel: cancel} return } func (db *conn) queryRow(c context.Context, query string, args ...interface{}) *Row { now := time.Now() defer slowLog(fmt.Sprintf("QueryRow addr: %s query(%s) args(%+v)", db.addr, query, args), now) t, ok := trace.FromContext(c) if ok { t = t.Fork(_family, "queryrow") t.SetTag(trace.String(trace.TagAddress, db.addr), trace.String(trace.TagComment, query)) } if err := db.breaker.Allow(); err != nil { _metricReqErr.Inc(db.addr, db.addr, "queryrow", "breaker") return &Row{db: db, t: t, err: err} } _, c, cancel := db.conf.QueryTimeout.Shrink(c) r := db.DB.QueryRowContext(c, query, args...) _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), db.addr, db.addr, "queryrow") return &Row{db: db, Row: r, query: query, args: args, t: t, cancel: cancel} } // Close closes the statement. func (s *Stmt) Close() (err error) { stmt, ok := s.stmt.Load().(*sql.Stmt) if ok { err = errors.WithStack(stmt.Close()) } return } func (s *Stmt) prepare() (st *sql.Stmt) { var ok bool if st, ok = s.stmt.Load().(*sql.Stmt); ok { return } var err error if st, err = s.db.Prepare(s.query); err == nil { s.stmt.Store(st) } return } // Exec executes a prepared statement with the given arguments and returns a // Result summarizing the effect of the statement. func (s *Stmt) Exec(c context.Context, args ...interface{}) (res sql.Result, err error) { now := time.Now() defer slowLog(fmt.Sprintf("Exec addr: %s query(%s) args(%+v)", s.db.addr, s.query, args), now) if s.tx { if s.t != nil { s.t.SetTag(trace.String(trace.TagAnnotation, s.query)) } } else if t, ok := trace.FromContext(c); ok { t = t.Fork(_family, "exec") t.SetTag(trace.String(trace.TagAddress, s.db.addr), trace.String(trace.TagComment, s.query)) defer t.Finish(&err) } if err = s.db.breaker.Allow(); err != nil { _metricReqErr.Inc(s.db.addr, s.db.addr, "stmt:exec", "breaker") return } stmt := s.prepare() if stmt == nil { err = ErrStmtNil return } _, c, cancel := s.db.conf.ExecTimeout.Shrink(c) res, err = stmt.ExecContext(c, args...) cancel() s.db.onBreaker(&err) _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), s.db.addr, s.db.addr, "stmt:exec") if err != nil { err = errors.Wrapf(err, "addr: %s exec:%s, args:%+v", s.db.addr, s.query, args) } return } // Query executes a prepared query statement with the given arguments and // returns the query results as a *Rows. func (s *Stmt) Query(c context.Context, args ...interface{}) (rows *Rows, err error) { now := time.Now() defer slowLog(fmt.Sprintf("Query addr: %s query(%s) args(%+v)", s.db.addr, s.query, args), now) if s.tx { if s.t != nil { s.t.SetTag(trace.String(trace.TagAnnotation, s.query)) } } else if t, ok := trace.FromContext(c); ok { t = t.Fork(_family, "query") t.SetTag(trace.String(trace.TagAddress, s.db.addr), trace.String(trace.TagComment, s.query)) defer t.Finish(&err) } if err = s.db.breaker.Allow(); err != nil { _metricReqErr.Inc(s.db.addr, s.db.addr, "stmt:query", "breaker") return } stmt := s.prepare() if stmt == nil { err = ErrStmtNil return } _, c, cancel := s.db.conf.QueryTimeout.Shrink(c) rs, err := stmt.QueryContext(c, args...) s.db.onBreaker(&err) _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), s.db.addr, s.db.addr, "stmt:query") if err != nil { err = errors.Wrapf(err, "addr: %s, query:%s, args:%+v", s.db.addr, s.query, args) cancel() return } rows = &Rows{Rows: rs, cancel: cancel} return } // QueryRow executes a prepared query statement with the given arguments. // If an error occurs during the execution of the statement, that error will // be returned by a call to Scan on the returned *Row, which is always non-nil. // If the query selects no rows, the *Row's Scan will return ErrNoRows. // Otherwise, the *Row's Scan scans the first selected row and discards the rest. func (s *Stmt) QueryRow(c context.Context, args ...interface{}) (row *Row) { now := time.Now() defer slowLog(fmt.Sprintf("QueryRow addr: %s query(%s) args(%+v)", s.db.addr, s.query, args), now) row = &Row{db: s.db, query: s.query, args: args} if s.tx { if s.t != nil { s.t.SetTag(trace.String(trace.TagAnnotation, s.query)) } } else if t, ok := trace.FromContext(c); ok { t = t.Fork(_family, "queryrow") t.SetTag(trace.String(trace.TagAddress, s.db.addr), trace.String(trace.TagComment, s.query)) row.t = t } if row.err = s.db.breaker.Allow(); row.err != nil { _metricReqErr.Inc(s.db.addr, s.db.addr, "stmt:queryrow", "breaker") return } stmt := s.prepare() if stmt == nil { return } _, c, cancel := s.db.conf.QueryTimeout.Shrink(c) row.Row = stmt.QueryRowContext(c, args...) row.cancel = cancel _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), s.db.addr, s.db.addr, "stmt:queryrow") return } func (s *Stmts) prepare(conn *conn) (st *Stmt) { if conn == nil { conn = s.db.conn() } s.mu.RLock() st = s.sts[conn.addr] s.mu.RUnlock() if st == nil { st = conn.prepared(s.query) s.mu.Lock() s.sts[conn.addr] = st s.mu.Unlock() } return } // Exec executes a prepared statement with the given arguments and returns a // Result summarizing the effect of the statement. func (s *Stmts) Exec(c context.Context, args ...interface{}) (res sql.Result, err error) { return s.prepare(nil).Exec(c, args...) } // Query executes a prepared query statement with the given arguments and // returns the query results as a *Rows. func (s *Stmts) Query(c context.Context, args ...interface{}) (rows *Rows, err error) { return s.prepare(nil).Query(c, args...) } // QueryRow executes a prepared query statement with the given arguments. // If an error occurs during the execution of the statement, that error will // be returned by a call to Scan on the returned *Row, which is always non-nil. // If the query selects no rows, the *Row's Scan will return ErrNoRows. // Otherwise, the *Row's Scan scans the first selected row and discards the rest. func (s *Stmts) QueryRow(c context.Context, args ...interface{}) (row *Row) { return s.prepare(nil).QueryRow(c, args...) } // Close closes the statement. func (s *Stmts) Close() (err error) { for _, st := range s.sts { if err = errors.WithStack(st.Close()); err != nil { return } } return } // Commit commits the transaction. func (tx *Tx) Commit() (err error) { err = tx.tx.Commit() tx.cancel() tx.db.onBreaker(&err) if tx.t != nil { tx.t.Finish(&err) } if err != nil { err = errors.WithStack(err) } return } // Rollback aborts the transaction. func (tx *Tx) Rollback() (err error) { err = tx.tx.Rollback() tx.cancel() tx.db.onBreaker(&err) if tx.t != nil { tx.t.Finish(&err) } if err != nil { err = errors.WithStack(err) } return } // Exec executes a query that doesn't return rows. For example: an INSERT and // UPDATE. func (tx *Tx) Exec(query string, args ...interface{}) (res sql.Result, err error) { now := time.Now() defer slowLog(fmt.Sprintf("Exec addr: %s query(%s) args(%+v)", tx.db.addr, query, args), now) if tx.t != nil { tx.t.SetTag(trace.String(trace.TagAnnotation, fmt.Sprintf("exec %s", query))) } res, err = tx.tx.ExecContext(tx.c, query, args...) _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), tx.db.addr, tx.db.addr, "tx:exec") if err != nil { err = errors.Wrapf(err, "addr: %s exec:%s, args:%+v", tx.db.addr, query, args) } return } // Query executes a query that returns rows, typically a SELECT. func (tx *Tx) Query(query string, args ...interface{}) (rows *Rows, err error) { if tx.t != nil { tx.t.SetTag(trace.String(trace.TagAnnotation, fmt.Sprintf("query %s", query))) } now := time.Now() defer slowLog(fmt.Sprintf("Query addr: %s query(%s) args(%+v)", tx.db.addr, query, args), now) defer func() { _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), tx.db.addr, tx.db.addr, "tx:query") }() rs, err := tx.tx.QueryContext(tx.c, query, args...) if err == nil { rows = &Rows{Rows: rs} } else { err = errors.Wrapf(err, "addr: %s, query:%s, args:%+v", tx.db.addr, query, args) } return } // QueryRow executes a query that is expected to return at most one row. // QueryRow always returns a non-nil value. Errors are deferred until Row's // Scan method is called. func (tx *Tx) QueryRow(query string, args ...interface{}) *Row { if tx.t != nil { tx.t.SetTag(trace.String(trace.TagAnnotation, fmt.Sprintf("queryrow %s", query))) } now := time.Now() defer slowLog(fmt.Sprintf("QueryRow addr: %s query(%s) args(%+v)", tx.db.addr, query, args), now) defer func() { _metricReqDur.Observe(int64(time.Since(now)/time.Millisecond), tx.db.addr, tx.db.addr, "tx:queryrow") }() r := tx.tx.QueryRowContext(tx.c, query, args...) return &Row{Row: r, db: tx.db, query: query, args: args} } // Stmt returns a transaction-specific prepared statement from an existing statement. func (tx *Tx) Stmt(stmt *Stmt) *Stmt { if stmt == nil { return nil } as, ok := stmt.stmt.Load().(*sql.Stmt) if !ok { return nil } ts := tx.tx.StmtContext(tx.c, as) st := &Stmt{query: stmt.query, tx: true, t: tx.t, db: tx.db} st.stmt.Store(ts) return st } // Stmts returns a transaction-specific prepared statement from an existing statement. func (tx *Tx) Stmts(stmt *Stmts) *Stmt { return tx.Stmt(stmt.prepare(tx.db)) } // Prepare creates a prepared statement for use within a transaction. // The returned statement operates within the transaction and can no longer be // used once the transaction has been committed or rolled back. // To use an existing prepared statement on this transaction, see Tx.Stmt. func (tx *Tx) Prepare(query string) (*Stmt, error) { if tx.t != nil { tx.t.SetTag(trace.String(trace.TagAnnotation, fmt.Sprintf("prepare %s", query))) } defer slowLog(fmt.Sprintf("Prepare addr: %s query(%s)", tx.db.addr, query), time.Now()) stmt, err := tx.tx.Prepare(query) if err != nil { err = errors.Wrapf(err, "addr: %s prepare %s", tx.db.addr, query) return nil, err } st := &Stmt{query: query, tx: true, t: tx.t, db: tx.db} st.stmt.Store(stmt) return st, nil } // parseDSNAddr parse dsn name and return addr. func parseDSNAddr(dsn string) (addr string) { if dsn == "" { return } cfg, err := mysql.ParseDSN(dsn) if err != nil { return } addr = cfg.Addr return } func genDSN(dsn, addr string) (res string) { cfg, err := mysql.ParseDSN(dsn) if err != nil { return } cfg.Addr = addr cfg.Net = "tcp" res = cfg.FormatDSN() return } func slowLog(statement string, now time.Time) { du := time.Since(now) if du > _slowLogDuration { log.Warn("%s slow log statement: %s time: %v", _family, statement, du) } }