From 2306d6180fc302124ffa23fb18783fc1fd794db9 Mon Sep 17 00:00:00 2001
From: Renan DelValle <rdelval@users.noreply.github.com>
Date: Mon, 22 Oct 2018 16:43:35 -0700
Subject: [PATCH] Adding force Implicit and force Explicit recon to gorealis.
 (#81)

---
 examples/client.go | 14 ++++++++++++++
 realis.go          | 37 +++++++++++++++++++++++++++++++++++++
 realis_e2e_test.go | 20 ++++++++++++++++++--
 3 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/examples/client.go b/examples/client.go
index ef59d77..cd6d726 100644
--- a/examples/client.go
+++ b/examples/client.go
@@ -605,6 +605,20 @@ func main() {
 			log.Fatalf("error: %+v\n", err.Error())
 		}
 
+	case "forceExplicitRecon":
+		fmt.Println("Force an explicit recon")
+		err := r.ForceExplicitTaskReconciliation(nil)
+		if err != nil {
+			log.Fatalf("error: %+v\n", err.Error())
+		}
+
+	case "forceImplicitRecon":
+		fmt.Println("Force an implicit recon")
+		err := r.ForceImplicitTaskReconciliation()
+		if err != nil {
+			log.Fatalf("error: %+v\n", err.Error())
+		}
+
 	default:
 		log.Fatal("Command not supported")
 	}
diff --git a/realis.go b/realis.go
index c028406..84c8b1e 100644
--- a/realis.go
+++ b/realis.go
@@ -78,6 +78,10 @@ type Realis interface {
 	GetQuota(role string) (*aurora.Response, error)
 	Snapshot() error
 	PerformBackup() error
+	// Force an Implicit reconciliation between Mesos and Aurora
+	ForceImplicitTaskReconciliation() error
+	// Force an Explicit reconciliation between Mesos and Aurora
+	ForceExplicitTaskReconciliation(batchSize *int32) error
 }
 
 type realisClient struct {
@@ -1112,3 +1116,36 @@ func (r *realisClient) PerformBackup() error {
 
 	return nil
 }
+
+func (r *realisClient) ForceImplicitTaskReconciliation() error {
+
+	_, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) {
+		return r.adminClient.TriggerImplicitTaskReconciliation()
+	})
+
+	if retryErr != nil {
+		return errors.Wrap(retryErr, "Unable to recover connection")
+	}
+
+	return nil
+}
+
+func (r *realisClient) ForceExplicitTaskReconciliation(batchSize *int32) error {
+
+	if batchSize != nil && *batchSize < 1 {
+		return errors.New("Invalid batch size.")
+	}
+	settings := aurora.NewExplicitReconciliationSettings()
+
+	settings.BatchSize = batchSize
+
+	_, retryErr := r.thriftCallWithRetries(func() (*aurora.Response, error) {
+		return r.adminClient.TriggerExplicitTaskReconciliation(settings)
+	})
+
+	if retryErr != nil {
+		return errors.Wrap(retryErr, "Unable to recover connection")
+	}
+
+	return nil
+}
diff --git a/realis_e2e_test.go b/realis_e2e_test.go
index 382fa53..3ea41ec 100644
--- a/realis_e2e_test.go
+++ b/realis_e2e_test.go
@@ -18,11 +18,10 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"sync"
 	"testing"
 	"time"
 
-	"sync"
-
 	"github.com/paypal/gorealis"
 	"github.com/paypal/gorealis/gen-go/apache/aurora"
 	"github.com/paypal/gorealis/response"
@@ -294,6 +293,7 @@ func TestRealisClient_CreateService(t *testing.T) {
 		IsService(true)
 
 	settings := realis.NewUpdateSettings()
+	settings.UpdateGroupSize = 2
 	job.InstanceCount(3)
 	resp, result, err := r.CreateService(job, settings)
 
@@ -521,3 +521,19 @@ func TestRealisClient_SetQuota(t *testing.T) {
 		fmt.Print("GetQuota Result", result.String())
 	})
 }
+
+func TestRealisClient_ForceImplicitTaskReconciliation(t *testing.T) {
+	err := r.ForceImplicitTaskReconciliation()
+	assert.NoError(t, err)
+}
+
+func TestRealisClient_ForceExplicitTaskReconciliation(t *testing.T) {
+	// Default value
+	err := r.ForceExplicitTaskReconciliation(nil)
+	assert.NoError(t, err)
+
+	// Custom batch value
+	var batchSize int32 = 32
+	err = r.ForceExplicitTaskReconciliation(&batchSize)
+	assert.NoError(t, err)
+}