From 553f4c8d13dc7eface08c1afe5aa72a4e8d9dea1 Mon Sep 17 00:00:00 2001
From: Shlomi Noach <shlomi-noach@github.com>
Date: Wed, 24 Aug 2016 11:39:44 +0200
Subject: [PATCH] concurrent row-count

---
 go/base/context.go    |  1 +
 go/cmd/gh-ost/main.go |  1 +
 go/logic/inspect.go   |  6 ++++--
 go/logic/migrator.go  | 28 +++++++++++++++++++++-------
 4 files changed, 27 insertions(+), 9 deletions(-)

diff --git a/go/base/context.go b/go/base/context.go
index dfc69bb..0ec67d1 100644
--- a/go/base/context.go
+++ b/go/base/context.go
@@ -49,6 +49,7 @@ type MigrationContext struct {
 	AlterStatement    string
 
 	CountTableRows           bool
+	ConcurrentCountTableRows bool
 	AllowedRunningOnMaster   bool
 	AllowedMasterMaster      bool
 	SwitchToRowBinlogFormat  bool
diff --git a/go/cmd/gh-ost/main.go b/go/cmd/gh-ost/main.go
index 512c537..78a4d5f 100644
--- a/go/cmd/gh-ost/main.go
+++ b/go/cmd/gh-ost/main.go
@@ -53,6 +53,7 @@ func main() {
 	flag.StringVar(&migrationContext.OriginalTableName, "table", "", "table name (mandatory)")
 	flag.StringVar(&migrationContext.AlterStatement, "alter", "", "alter statement (mandatory)")
 	flag.BoolVar(&migrationContext.CountTableRows, "exact-rowcount", false, "actually count table rows as opposed to estimate them (results in more accurate progress estimation)")
+	flag.BoolVar(&migrationContext.ConcurrentCountTableRows, "concurrent-rowcount", false, "(with --exact-rowcount), when true: count rows after row-copy begins, concurrently, and adjust row estimate later on; defaults false: first count rows, then start row copy")
 	flag.BoolVar(&migrationContext.AllowedRunningOnMaster, "allow-on-master", false, "allow this migration to run directly on master. Preferably it would run on a replica")
 	flag.BoolVar(&migrationContext.AllowedMasterMaster, "allow-master-master", false, "explicitly allow running in a master-master setup")
 	flag.BoolVar(&migrationContext.NullableUniqueKeyAllowed, "allow-nullable-unique-key", false, "allow gh-ost to migrate based on a unique key with nullable columns. As long as no NULL values exist, this should be OK. If NULL values exist in chosen key, data may be corrupted. Use at your own risk!")
diff --git a/go/logic/inspect.go b/go/logic/inspect.go
index 8377a9c..644613d 100644
--- a/go/logic/inspect.go
+++ b/go/logic/inspect.go
@@ -440,12 +440,14 @@ func (this *Inspector) CountTableRows() error {
 	log.Infof("As instructed, I'm issuing a SELECT COUNT(*) on the table. This may take a while")
 
 	query := fmt.Sprintf(`select /* gh-ost */ count(*) as rows from %s.%s`, sql.EscapeName(this.migrationContext.DatabaseName), sql.EscapeName(this.migrationContext.OriginalTableName))
-	if err := this.db.QueryRow(query).Scan(&this.migrationContext.RowsEstimate); err != nil {
+	var rowsEstimate int64
+	if err := this.db.QueryRow(query).Scan(&rowsEstimate); err != nil {
 		return err
 	}
+	atomic.StoreInt64(&this.migrationContext.RowsEstimate, rowsEstimate)
 	this.migrationContext.UsedRowsEstimateMethod = base.CountRowsEstimate
 
-	log.Infof("Exact number of rows via COUNT: %d", this.migrationContext.RowsEstimate)
+	log.Infof("Exact number of rows via COUNT: %d", rowsEstimate)
 
 	return nil
 }
diff --git a/go/logic/migrator.go b/go/logic/migrator.go
index 65d869a..c7e805e 100644
--- a/go/logic/migrator.go
+++ b/go/logic/migrator.go
@@ -361,6 +361,24 @@ func (this *Migrator) validateStatement() (err error) {
 	return nil
 }
 
+func (this *Migrator) countTableRows() (err error) {
+	if !this.migrationContext.CountTableRows {
+		// Not counting; we stay with an estimate
+		return nil
+	}
+	if this.migrationContext.Noop {
+		log.Debugf("Noop operation; not really counting table rows")
+		return nil
+	}
+	if this.migrationContext.ConcurrentCountTableRows {
+		go this.inspector.CountTableRows()
+		log.Infof("As instructed, counting rows in the background; meanwhile I will use an estimated count, and will update it later on")
+		// and we ignore errors, because this turns to be a background job
+		return nil
+	}
+	return this.inspector.CountTableRows()
+}
+
 // Migrate executes the complete migration logic. This is *the* major gh-ost function.
 func (this *Migrator) Migrate() (err error) {
 	log.Infof("Migrating %s.%s", sql.EscapeName(this.migrationContext.DatabaseName), sql.EscapeName(this.migrationContext.OriginalTableName))
@@ -402,12 +420,8 @@ func (this *Migrator) Migrate() (err error) {
 	}
 	defer this.server.RemoveSocketFile()
 
-	if this.migrationContext.CountTableRows {
-		if this.migrationContext.Noop {
-			log.Debugf("Noop operation; not really counting table rows")
-		} else if err := this.inspector.CountTableRows(); err != nil {
-			return err
-		}
+	if err := this.countTableRows(); err != nil {
+		return err
 	}
 
 	if err := this.addDMLEventsListener(); err != nil {
@@ -970,7 +984,7 @@ func (this *Migrator) printStatus(rule PrintStatusRule, writers ...io.Writer) {
 
 	var etaSeconds float64 = math.MaxFloat64
 	eta := "N/A"
-	if atomic.LoadInt64(&this.migrationContext.CountingRowsFlag) > 0 {
+	if atomic.LoadInt64(&this.migrationContext.CountingRowsFlag) > 0 && !this.migrationContext.ConcurrentCountTableRows {
 		eta = "counting rows"
 	} else if atomic.LoadInt64(&this.migrationContext.IsPostponingCutOver) > 0 {
 		eta = "postponing cut-over"