Implement Pinger interface

AUTHORS

@@ -48,6 +48,7 @@ Runrioter Wung <runrioter at gmail.com>
 Soroush Pour <me at soroushjp.com>
 Stan Putrya <root.vagner at gmail.com>
 Stanley Gunawan <gunawan.stanley at gmail.com>
+Steven Huang <photon3108 at gmail.com>
 Xiangyu Hu <xiangyu.hu at outlook.com>
 Xiaobing Jiang <s7v7nislands at gmail.com>
 Xiuming Chen <cc at cxm.cc>

buffer.go

@@ -22,18 +22,25 @@ const defaultBufSize = 4096
 // The buffer is similar to bufio.Reader / Writer but zero-copy-ish
 // Also highly optimized for this particular use case.
 type buffer struct {
-	buf     []byte
+	buf    []byte
+	idx    int
+	length int
+
 	nc      net.Conn
-	idx     int
-	length  int
 	timeout time.Duration
+	ctx     Context
 }
 
 func newBuffer(nc net.Conn) buffer {
+	return newBufferContext(nc, nil)
+}
+
+func newBufferContext(nc net.Conn, ctx Context) buffer {
 	var b [defaultBufSize]byte
 	return buffer{
 		buf: b[:],
 		nc:  nc,
+		ctx: ctx,
 	}
 }
 
@@ -60,7 +67,15 @@ func (b *buffer) fill(need int) error {
 
 	for {
 		if b.timeout > 0 {
-			if err := b.nc.SetReadDeadline(time.Now().Add(b.timeout)); err != nil {
+			deadline := time.Now().Add(b.timeout)
+			if b.ctx != nil {
+				if ctxDeadline, ok := b.ctx.Deadline(); ok {
+					if !ctxDeadline.IsZero() && ctxDeadline.Before(deadline) {
+						deadline = ctxDeadline
+					}
+				}
+			}
+			if err := b.nc.SetReadDeadline(deadline); err != nil {
 				return err
 			}
 		}

connection.go

@@ -17,6 +17,11 @@ import (
 	"time"
 )
 
+// Define mysql.Context for the compatibility before go1.7.
+type Context interface {
+	Deadline() (deadline time.Time, ok bool)
+}
+
 type mysqlConn struct {
 	buf              buffer
 	netConn          net.Conn
@@ -31,6 +36,7 @@ type mysqlConn struct {
 	sequence         uint8
 	parseTime        bool
 	strict           bool
+	ctx              Context
 }
 
 // Handles parameters set in DSN after the connection is established
@@ -389,3 +395,16 @@ func (mc *mysqlConn) getSystemVar(name string) ([]byte, error) {
 	}
 	return nil, err
 }
+
+// Make all others get the context from the mysqlConn instead of copying it as
+// a member of struct everywhere.
+type ContextWrapper struct {
+	mc *mysqlConn
+}
+
+func (wrapper *ContextWrapper) Deadline() (time.Time, bool) {
+	if wrapper.mc.ctx == nil {
+		return time.Time{}, false
+	}
+	return wrapper.mc.ctx.Deadline()
+}

connection_go18.go

@@ -0,0 +1,90 @@
+// +build go1.8
+
+// Go MySQL Driver - A MySQL-Driver for Go's database/sql package
+//
+// Copyright 2017 The Go-MySQL-Driver Authors. All rights reserved.
+//
+// This Source Code Form is subject to the terms of the Mozilla Public
+// License, v. 2.0. If a copy of the MPL was not distributed with this file,
+// You can obtain one at http://mozilla.org/MPL/2.0/.
+
+package mysql
+
+import (
+	"context"
+	"database/sql/driver"
+	"net"
+	"time"
+)
+
+func (mc *mysqlConn) Ping(ctx context.Context) error {
+	var err error
+
+	if ctx == nil {
+		return mc.pingImpl()
+	}
+	mc.ctx = ctx
+	defer func() {
+		mc.ctx = nil
+	}()
+
+	if deadline, ok := ctx.Deadline(); ok {
+		if err = mc.netConn.SetDeadline(deadline); err != nil {
+			return err
+		}
+	}
+
+	if ctx.Done() == nil {
+		return mc.pingImpl()
+	}
+
+	result := make(chan error)
+	go func() {
+		result <- mc.pingImpl()
+	}()
+
+	select {
+	case <-ctx.Done():
+		// Because buffer.fill() and mysqlConn.writePacket() may overwrite the
+		// deadline of read/write again and again in the above goroutine,
+		// it has to use a loop to make sure SetDeadline() works.
+	UpdateDeadlineLoop:
+		for {
+			// Copy it as a local variable because mc.netConn may be closed and
+			// assigned nil in the above goroutine.
+			netConn := mc.netConn
+			if netConn != nil {
+				errDeadline := netConn.SetDeadline(time.Now())
+				if errDeadline != nil {
+					errLog.Print(errDeadline)
+				}
+			}
+			select {
+			case <-time.After(200 * time.Millisecond):
+				// Prevent from leaking the above goroutine.
+			case err = <-result:
+				break UpdateDeadlineLoop
+			}
+		}
+	case err = <-result:
+	}
+
+	if netErr, ok := err.(net.Error); ok {
+		// We don't know where it timed out and it may leave some redundant data
+		// in the connection so make it to be closed by DB.puConn() of the
+		// caller.
+		if netErr.Timeout() {
+			return driver.ErrBadConn
+		}
+	}
+	return err
+}
+
+func (mc *mysqlConn) pingImpl() error {
+	if err := mc.writeCommandPacket(comPing); err != nil {
+		return err
+	}
+
+	_, err := mc.readResultOK()
+	return err
+}