From c91b2d098d66682dbb17a770f1f0f8504fa3e7cb Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Tue, 29 Mar 2016 19:24:44 -0700
Subject: [PATCH] clientv3: AlarmList and AlarmDisarm

---
 clientv3/client.go      |  2 +-
 clientv3/maintenance.go | 95 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 96 insertions(+), 1 deletion(-)

diff --git a/clientv3/client.go b/clientv3/client.go
index cec5695b2..95a603a5d 100644
--- a/clientv3/client.go
+++ b/clientv3/client.go
@@ -184,7 +184,7 @@ func newClient(cfg *Config) (*Client, error) {
 	client.Lease = NewLease(client)
 	client.Watcher = NewWatcher(client)
 	client.Auth = NewAuth(client)
-	client.Maintenance = &maintenance{c: client}
+	client.Maintenance = NewMaintenance(client)
 	if cfg.Logger != nil {
 		logger.Set(cfg.Logger)
 	} else {
diff --git a/clientv3/maintenance.go b/clientv3/maintenance.go
index 158293ebd..6db7b3e81 100644
--- a/clientv3/maintenance.go
+++ b/clientv3/maintenance.go
@@ -15,15 +15,26 @@
 package clientv3
 
 import (
+	"sync"
+
 	pb "github.com/coreos/etcd/etcdserver/etcdserverpb"
 	"golang.org/x/net/context"
+	"google.golang.org/grpc"
 )
 
 type (
 	DefragmentResponse pb.DefragmentResponse
+	AlarmResponse      pb.AlarmResponse
+	AlarmMember        pb.AlarmMember
 )
 
 type Maintenance interface {
+	// AlarmList gets all active alarms.
+	AlarmList(ctx context.Context) (*AlarmResponse, error)
+
+	// AlarmDisarm disarms a given alarm.
+	AlarmDisarm(ctx context.Context, m *AlarmMember) (*AlarmResponse, error)
+
 	// Defragment defragments storage backend of the etcd member with given endpoint.
 	// Defragment is only needed when deleting a large number of keys and want to reclaim
 	// the resources.
@@ -36,6 +47,72 @@ type Maintenance interface {
 
 type maintenance struct {
 	c *Client
+
+	mu     sync.Mutex
+	conn   *grpc.ClientConn // conn in-use
+	remote pb.MaintenanceClient
+}
+
+func NewMaintenance(c *Client) Maintenance {
+	conn := c.ActiveConnection()
+	return &maintenance{
+		c:      c,
+		conn:   conn,
+		remote: pb.NewMaintenanceClient(conn),
+	}
+}
+
+func (m *maintenance) AlarmList(ctx context.Context) (*AlarmResponse, error) {
+	req := &pb.AlarmRequest{
+		Action:   pb.AlarmRequest_GET,
+		MemberID: 0,                 // all
+		Alarm:    pb.AlarmType_NONE, // all
+	}
+	for {
+		resp, err := m.getRemote().Alarm(ctx, req)
+		if err == nil {
+			return (*AlarmResponse)(resp), nil
+		}
+		if isHalted(ctx, err) {
+			return nil, err
+		}
+		if err = m.switchRemote(err); err != nil {
+			return nil, err
+		}
+	}
+}
+
+func (m *maintenance) AlarmDisarm(ctx context.Context, am *AlarmMember) (*AlarmResponse, error) {
+	req := &pb.AlarmRequest{
+		Action:   pb.AlarmRequest_DEACTIVATE,
+		MemberID: am.MemberID,
+		Alarm:    am.Alarm,
+	}
+
+	if req.MemberID == 0 && req.Alarm == pb.AlarmType_NONE {
+		ar, err := m.AlarmList(ctx)
+		if err != nil {
+			return nil, err
+		}
+		ret := AlarmResponse{}
+		for _, am := range ar.Alarms {
+			dresp, derr := m.AlarmDisarm(ctx, (*AlarmMember)(am))
+			if derr != nil {
+				return nil, derr
+			}
+			ret.Alarms = append(ret.Alarms, dresp.Alarms...)
+		}
+		return &ret, nil
+	}
+
+	resp, err := m.getRemote().Alarm(ctx, req)
+	if err == nil {
+		return (*AlarmResponse)(resp), nil
+	}
+	if !isHalted(ctx, err) {
+		go m.switchRemote(err)
+	}
+	return nil, err
 }
 
 func (m *maintenance) Defragment(ctx context.Context, endpoint string) (*DefragmentResponse, error) {
@@ -50,3 +127,21 @@ func (m *maintenance) Defragment(ctx context.Context, endpoint string) (*Defragm
 	}
 	return (*DefragmentResponse)(resp), nil
 }
+
+func (m *maintenance) getRemote() pb.MaintenanceClient {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	return m.remote
+}
+
+func (m *maintenance) switchRemote(prevErr error) error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	newConn, err := m.c.retryConnection(m.conn, prevErr)
+	if err != nil {
+		return err
+	}
+	m.conn = newConn
+	m.remote = pb.NewMaintenanceClient(m.conn)
+	return nil
+}