From 46bbea2a321c234fbf570dca98e776a7f8d5cc40 Mon Sep 17 00:00:00 2001
From: Shlomi Noach <shlomi-noach@github.com>
Date: Fri, 29 Jul 2016 10:40:23 +0200
Subject: [PATCH] ETA counting rows, fixed copy time on count

---
 go/base/context.go   | 13 +++++++++++++
 go/logic/inspect.go  |  7 +++++++
 go/logic/migrator.go |  6 ++++--
 3 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/go/base/context.go b/go/base/context.go
index f906aeb..103ce43 100644
--- a/go/base/context.go
+++ b/go/base/context.go
@@ -111,6 +111,7 @@ type MigrationContext struct {
 	throttleReason            string
 	throttleMutex             *sync.Mutex
 	IsPostponingCutOver       int64
+	CountingRowsFlag          int64
 
 	OriginalTableColumns             *sql.ColumnList
 	OriginalTableUniqueKeys          [](*sql.UniqueKey)
@@ -261,11 +262,23 @@ func (this *MigrationContext) ElapsedTime() time.Duration {
 	return time.Now().Sub(this.StartTime)
 }
 
+// MarkRowCopyStartTime
+func (this *MigrationContext) MarkRowCopyStartTime() {
+	this.throttleMutex.Lock()
+	defer this.throttleMutex.Unlock()
+	this.RowCopyStartTime = time.Now()
+}
+
 // ElapsedRowCopyTime returns time since starting to copy chunks of rows
 func (this *MigrationContext) ElapsedRowCopyTime() time.Duration {
 	this.throttleMutex.Lock()
 	defer this.throttleMutex.Unlock()
 
+	if this.RowCopyStartTime.IsZero() {
+		// Row copy hasn't started yet
+		return 0
+	}
+
 	if this.RowCopyEndTime.IsZero() {
 		return time.Now().Sub(this.RowCopyStartTime)
 	}
diff --git a/go/logic/inspect.go b/go/logic/inspect.go
index 81ee528..98b5829 100644
--- a/go/logic/inspect.go
+++ b/go/logic/inspect.go
@@ -9,6 +9,7 @@ import (
 	gosql "database/sql"
 	"fmt"
 	"strings"
+	"sync/atomic"
 
 	"github.com/github/gh-ost/go/base"
 	"github.com/github/gh-ost/go/mysql"
@@ -364,13 +365,19 @@ func (this *Inspector) estimateTableRowsViaExplain() error {
 
 // CountTableRows counts exact number of rows on the original table
 func (this *Inspector) CountTableRows() error {
+	atomic.StoreInt64(&this.migrationContext.CountingRowsFlag, 1)
+	defer atomic.StoreInt64(&this.migrationContext.CountingRowsFlag, 0)
+
 	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 {
 		return err
 	}
 	this.migrationContext.UsedRowsEstimateMethod = base.CountRowsEstimate
+
 	log.Infof("Exact number of rows via COUNT: %d", this.migrationContext.RowsEstimate)
+
 	return nil
 }
 
diff --git a/go/logic/migrator.go b/go/logic/migrator.go
index 607a2fd..54c5a13 100644
--- a/go/logic/migrator.go
+++ b/go/logic/migrator.go
@@ -418,7 +418,7 @@ func (this *Migrator) Migrate() (err error) {
 	go this.initiateThrottler()
 	go this.executeWriteFuncs()
 	go this.iterateChunks()
-	this.migrationContext.RowCopyStartTime = time.Now()
+	this.migrationContext.MarkRowCopyStartTime()
 	go this.initiateStatus()
 
 	log.Debugf("Operating until row copy is complete")
@@ -951,7 +951,9 @@ func (this *Migrator) printStatus(rule PrintStatusRule, writers ...io.Writer) {
 
 	var etaSeconds float64 = math.MaxFloat64
 	eta := "N/A"
-	if atomic.LoadInt64(&this.migrationContext.IsPostponingCutOver) > 0 {
+	if atomic.LoadInt64(&this.migrationContext.CountingRowsFlag) > 0 {
+		eta = "counting rows"
+	} else if atomic.LoadInt64(&this.migrationContext.IsPostponingCutOver) > 0 {
 		eta = "postponing cut-over"
 	} else if isThrottled, throttleReason := this.migrationContext.IsThrottled(); isThrottled {
 		eta = fmt.Sprintf("throttled, %s", throttleReason)