From 9e166fae542dfb5ae916d267a35cffed52d1494a Mon Sep 17 00:00:00 2001
From: lenhattan86 <lenhattan86@users.noreply.github.com>
Date: Wed, 14 Jul 2021 11:19:54 -0700
Subject: [PATCH] add SlaRestartInstances api

---
 jobUpdate.go |  5 +++++
 realis.go    | 44 ++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+)

diff --git a/jobUpdate.go b/jobUpdate.go
index 3ef2cb4..5ccbd1a 100644
--- a/jobUpdate.go
+++ b/jobUpdate.go
@@ -284,3 +284,8 @@ func (j *JobUpdate) PartitionPolicy(reschedule bool, delay int64) *JobUpdate {
 	})
 	return j
 }
+
+func (j *JobUpdate) SlaPolicy(sla *aurora.SlaPolicy) *JobUpdate {
+	j.task.task.SlaPolicy = sla
+	return j
+}
diff --git a/realis.go b/realis.go
index a7c536e..bc9171a 100644
--- a/realis.go
+++ b/realis.go
@@ -522,6 +522,50 @@ func (c *Client) RestartInstances(key aurora.JobKey, instances ...int32) error {
 	return nil
 }
 
+// Restarts specific instances specified with slaAware.
+// If slaPolicy is nil, uses the existing slaPolicy of taskConfig.
+func (c *Client) SlaRestartInstances(slaPolicy *aurora.SlaPolicy, jobKey *aurora.JobKey, instances ...int32) error {
+	c.logger.DebugPrintf("SlaRestartInstances Thrift Payload: %v %+v %v\n", slaPolicy, jobKey, instances)
+
+	if len(instances) == 0 {
+		c.logger.DebugPrintf("it is not necessary to restart 0 instances")
+		return nil
+	}
+
+	jobSummary, err := c.GetJobSummary(jobKey.Role)
+	if err != nil {
+		return err
+	}
+	var jobConfig *aurora.JobConfiguration
+	for _, s := range jobSummary.Summaries {
+		if s.Job.Key.Environment == jobKey.Environment && s.Job.Key.Name == jobKey.Name {
+			jobConfig = s.Job
+		}
+	}
+	if jobConfig == nil {
+		return fmt.Errorf("failed to find %v", jobKey)
+	}
+
+	// create job update request
+	jobUpdate := JobUpdateFromConfig(jobConfig.TaskConfig)
+	jobUpdate.
+		SlaAware(true).
+		InstanceCount(jobConfig.InstanceCount)
+	if slaPolicy != nil {
+		jobUpdate.SlaPolicy(slaPolicy)
+	}
+	for _, v := range instances {
+		jobUpdate.AddInstanceRange(v, v)
+	}
+
+	msg := fmt.Sprintf("SlaRestartInstances %v-%v via StartJobUpdate", jobKey, instances)
+	if _, err := c.StartJobUpdate(jobUpdate, ""); err != nil {
+		return errors.Wrap(err, msg)
+	}
+
+	return nil
+}
+
 // Restarts all active tasks under a job configuration.
 func (c *Client) RestartJob(key aurora.JobKey) error {