Update example

README.md

@@ -19,7 +19,7 @@ import (
 func main() {
 	stepsCount := 1000
 
-	eta := eta.New(time.Minute, stepsCount)
+	eta := eta.New(stepsCount)
 
 	processed := 0
 
@@ -39,6 +39,5 @@ func main() {
 		time.Sleep(time.Second) // Update progress every second
 		fmt.Fprintf(os.Stderr, "\rProcessed %d of %d, ETA: %s", processed, stepsCount, eta.Eta().Format("15:04:05"))
 	}
-
 }
 ```

consts.go

@@ -0,0 +1,8 @@
+package eta
+
+import "time"
+
+const (
+	defaultPeriodCount    = 10
+	defaultPeriodDuration = time.Minute
+)

eta.go

@@ -13,6 +13,9 @@ type Calculator struct {
 	// Expected processing count
 	TotalCount int
 
+	// Number of periods to store
+	PeriodCount int
+
 	periodDuration   time.Duration
 	currentPeriod    time.Time
 	currentProcessed int
@@ -22,12 +25,18 @@ type Calculator struct {
 }
 
 // New return new ETA calculator
-func New(periodDuration time.Duration, totalCount int) *Calculator {
+func New(totalCount int) *Calculator {
+	return NewCustom(totalCount, defaultPeriodDuration)
+}
+
+// NewCustom return new ETA calculator with custom params
+func NewCustom(totalCount int, periodDuration time.Duration) *Calculator {
 	now := time.Now()
 
 	etaCalc := &Calculator{
 		startTime:      now,
 		TotalCount:     totalCount,
+		PeriodCount:    defaultPeriodCount,
 		currentPeriod:  now.Truncate(periodDuration),
 		periodDuration: periodDuration}
 
@@ -59,8 +68,8 @@ func (ec *Calculator) Increment(n int) {
 		ec.currentPeriod = period
 	}
 
-	if len(ec.stats) > 10 {
-		ec.stats = ec.stats[:10]
+	if len(ec.stats) > ec.PeriodCount {
+		ec.stats = ec.stats[:ec.PeriodCount]
 	}
 }
 
@@ -78,33 +87,15 @@ func (ec *Calculator) Last() time.Time {
 	return time.Now().Add(lastPeriodSpeed * time.Duration(ec.TotalCount-ec.processed))
 }
 
-// Eta returns ETA based on total time and total processed items count
-func (ec *Calculator) Eta() time.Time {
-	if ec.processed == 0 {
-		return time.Time{}
-	}
+// cycleTime returns cycle time based on total time and total processed items count
+func (ec *Calculator) cycleTime(now time.Time) time.Duration {
+	elapsedTime := time.Since(ec.startTime)
 
-	ec.mu.RLock()
-	defer ec.mu.RUnlock()
-
-	now := time.Now()
-	elapsedTime := now.Sub(ec.startTime)
-	avgSpeed := elapsedTime / time.Duration(ec.processed)
-
-	return now.Add(avgSpeed * time.Duration(ec.TotalCount-ec.processed))
+	return elapsedTime / time.Duration(ec.processed)
 }
 
-// Average returns ETA based on average processing speed of last periods
-func (ec *Calculator) Average() time.Time {
-	if len(ec.stats) == 0 {
-		return ec.Eta()
-	}
-
-	ec.mu.RLock()
-	defer ec.mu.RUnlock()
-
-	now := time.Now()
-
+// averageCycleTime returns cycle time based on average processing speed of last periods
+func (ec *Calculator) averageCycleTime() time.Duration {
 	processed := ec.stats[len(ec.stats)-1]
 	startPeriod := ec.currentPeriod.Add(-ec.periodDuration)
 
@@ -114,25 +105,14 @@ func (ec *Calculator) Average() time.Time {
 	}
 
 	if processed == 0 {
-		return time.Time{}
+		return time.Duration(0)
 	}
 
-	avgSpeed := ec.currentPeriod.Sub(startPeriod) / time.Duration(processed)
-
-	return now.Add(time.Duration(ec.TotalCount-ec.processed) * avgSpeed)
+	return ec.currentPeriod.Sub(startPeriod) / time.Duration(processed)
 }
 
-// Optimistic returns ETA based on detected maximum of processing speed
-func (ec *Calculator) Optimistic() time.Time {
-	if len(ec.stats) == 0 {
-		return ec.Eta()
-	}
-
-	ec.mu.RLock()
-	defer ec.mu.RUnlock()
-
-	now := time.Now()
-
+// optimisticCycleTime returns cycle time based on detected maximum of processing speed
+func (ec *Calculator) optimisticCycleTime() time.Duration {
 	var maxSpeed time.Duration
 	if ec.stats[len(ec.stats)-1] > 0 {
 		maxSpeed = ec.periodDuration / time.Duration(ec.stats[len(ec.stats)-1])
@@ -151,20 +131,11 @@ func (ec *Calculator) Optimistic() time.Time {
 		}
 	}
 
-	return now.Add(time.Duration(ec.TotalCount-ec.processed) * maxSpeed)
+	return maxSpeed
 }
 
-// Pessimistic returns ETA based on detected minimum of processing speed
-func (ec *Calculator) Pessimistic() time.Time {
-	if len(ec.stats) == 0 {
-		return ec.Eta()
-	}
-
-	ec.mu.RLock()
-	defer ec.mu.RUnlock()
-
-	now := time.Now()
-
+// pessimisticCycleTime returns cycle time based on detected minimum of processing speed
+func (ec *Calculator) pessimisticCycleTime() time.Duration {
 	var minSpeed time.Duration
 	if ec.stats[len(ec.stats)-1] > 0 {
 		minSpeed = ec.periodDuration / time.Duration(ec.stats[len(ec.stats)-1])
@@ -186,5 +157,71 @@ func (ec *Calculator) Pessimistic() time.Time {
 		}
 	}
 
-	return now.Add(time.Duration(ec.TotalCount-ec.processed) * minSpeed * time.Duration(1+nulPeriods))
+	return minSpeed * time.Duration(1+nulPeriods)
+}
+
+// Eta returns ETA based on total time and total processed items count
+func (ec *Calculator) Eta() time.Time {
+	if ec.processed == 0 {
+		return time.Time{}
+	}
+
+	ec.mu.RLock()
+	defer ec.mu.RUnlock()
+
+	now := time.Now()
+	avgCycleTime := ec.cycleTime(now)
+
+	return now.Add(avgCycleTime * time.Duration(ec.TotalCount-ec.processed))
+}
+
+// Average returns ETA based on average processing speed of last periods
+func (ec *Calculator) Average() time.Time {
+	if len(ec.stats) == 0 {
+		return ec.Eta()
+	}
+
+	ec.mu.RLock()
+	defer ec.mu.RUnlock()
+
+	avgCycleTime := ec.averageCycleTime()
+	if avgCycleTime == 0 {
+		return time.Time{}
+	}
+
+	return time.Now().Add(time.Duration(ec.TotalCount-ec.processed) * avgCycleTime)
+}
+
+// Optimistic returns ETA based on detected maximum of processing speed
+func (ec *Calculator) Optimistic() time.Time {
+	if len(ec.stats) == 0 {
+		return ec.Eta()
+	}
+
+	ec.mu.RLock()
+	defer ec.mu.RUnlock()
+
+	optimisticCycleTime := ec.optimisticCycleTime()
+	if optimisticCycleTime == 0 {
+		return time.Time{}
+	}
+
+	return time.Now().Add(time.Duration(ec.TotalCount-ec.processed) * ec.optimisticCycleTime())
+}
+
+// Pessimistic returns ETA based on detected minimum of processing speed
+func (ec *Calculator) Pessimistic() time.Time {
+	if len(ec.stats) == 0 {
+		return ec.Eta()
+	}
+
+	ec.mu.RLock()
+	defer ec.mu.RUnlock()
+
+	pessimisticCycleTime := ec.pessimisticCycleTime()
+	if pessimisticCycleTime == 0 {
+		return time.Time{}
+	}
+
+	return time.Now().Add(time.Duration(ec.TotalCount-ec.processed) * pessimisticCycleTime)
 }