From 3b780e52045f86cc943ebd1c67ead7b3d6f313d6 Mon Sep 17 00:00:00 2001
From: Iwasaki Yudai <yudai.iwasaki@ntti3.com>
Date: Fri, 16 Jun 2017 17:11:48 -0700
Subject: [PATCH] Add new option for revision based auto compaction

---
 compactor/compactor.go | 22 ++++++++++++++++++++++
 embed/config.go        |  1 +
 embed/etcd.go          |  1 +
 etcdmain/config.go     |  3 ++-
 etcdmain/help.go       |  4 +++-
 etcdserver/config.go   |  1 +
 etcdserver/server.go   |  9 ++++++---
 7 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/compactor/compactor.go b/compactor/compactor.go
index 723456c066a9..63c0d425bb5b 100644
--- a/compactor/compactor.go
+++ b/compactor/compactor.go
@@ -15,6 +15,7 @@
 package compactor
 
 import (
+	"fmt"
 	"time"
 
 	"github.com/coreos/pkg/capnslog"
@@ -30,8 +31,18 @@ var (
 const (
 	checkCompactionInterval   = 5 * time.Minute
 	executeCompactionInterval = time.Hour
+
+	ModePeriodic   = "period"
+	ModeRevisional = "revision"
 )
 
+type Compactor interface {
+	Run()
+	Stop()
+	Pause()
+	Resume()
+}
+
 type Compactable interface {
 	Compact(ctx context.Context, r *pb.CompactionRequest) (*pb.CompactionResponse, error)
 }
@@ -39,3 +50,14 @@ type Compactable interface {
 type RevGetter interface {
 	Rev() int64
 }
+
+func New(mode string, retention int, rg RevGetter, c Compactable) (Compactor, error) {
+	switch mode {
+	case ModePeriodic:
+		return NewPeriodic(retention, rg, c), nil
+	case ModeRevisional:
+		return NewRevisional(int64(retention), rg, c), nil
+	default:
+		return nil, fmt.Errorf("unsupported compaction mode %s", mode)
+	}
+}
diff --git a/embed/config.go b/embed/config.go
index 44acc2bf2c77..87a451262e2a 100644
--- a/embed/config.go
+++ b/embed/config.go
@@ -80,6 +80,7 @@ type Config struct {
 	Name                    string `json:"name"`
 	SnapCount               uint64 `json:"snapshot-count"`
 	AutoCompactionRetention int    `json:"auto-compaction-retention"`
+	AutoCompactionMode      string    `json:"auto-compaction-mode`
 
 	// TickMs is the number of milliseconds between heartbeat ticks.
 	// TODO: decouple tickMs and heartbeat tick (current heartbeat tick = 1).
diff --git a/embed/etcd.go b/embed/etcd.go
index b8e170f0681a..2105dc3debe8 100644
--- a/embed/etcd.go
+++ b/embed/etcd.go
@@ -138,6 +138,7 @@ func StartEtcd(inCfg *Config) (e *Etcd, err error) {
 		TickMs:                  cfg.TickMs,
 		ElectionTicks:           cfg.ElectionTicks(),
 		AutoCompactionRetention: cfg.AutoCompactionRetention,
+		AutoCompactionMode:      cfg.AutoCompactionMode,
 		QuotaBackendBytes:       cfg.QuotaBackendBytes,
 		MaxTxnOps:               cfg.MaxTxnOps,
 		MaxRequestBytes:         cfg.MaxRequestBytes,
diff --git a/etcdmain/config.go b/etcdmain/config.go
index 8f257e0ef9eb..037de5450199 100644
--- a/etcdmain/config.go
+++ b/etcdmain/config.go
@@ -199,7 +199,8 @@ func newConfig() *config {
 	// version
 	fs.BoolVar(&cfg.printVersion, "version", false, "Print the version and exit.")
 
-	fs.IntVar(&cfg.AutoCompactionRetention, "auto-compaction-retention", 0, "Auto compaction retention for mvcc key value store in hour. 0 means disable auto compaction.")
+	fs.IntVar(&cfg.AutoCompactionRetention, "auto-compaction-retention", 0, "Auto compaction retention for mvcc key value store. 0 means disable auto compaction.")
+	fs.StringVar(&cfg.AutoCompactionMode, "auto-compaction-mode", "period", "Interpreted 'auto-comapction-retention' as hours when 'period', as revision numbers when 'revision'.")
 
 	// pprof profiler via HTTP
 	fs.BoolVar(&cfg.EnablePprof, "enable-pprof", false, "Enable runtime profiling data via HTTP server. Address is at client URL + \"/debug/pprof/\"")
diff --git a/etcdmain/help.go b/etcdmain/help.go
index cbead2e80cdb..767dedcf8313 100644
--- a/etcdmain/help.go
+++ b/etcdmain/help.go
@@ -97,7 +97,9 @@ clustering flags:
 	--strict-reconfig-check
 		reject reconfiguration requests that would cause quorum loss.
 	--auto-compaction-retention '0'
-		auto compaction retention in hour. 0 means disable auto compaction.
+		auto compaction retention length. 0 means disable auto compaction.
+	--auto-compaction-mode 'period'
+		'period' means hours, 'revision' means revision numbers to retain by auto compaction
 	--enable-v2
 		Accept etcd V2 client requests.
 
diff --git a/etcdserver/config.go b/etcdserver/config.go
index 7ab77d986c3b..f6ed1f1bae7a 100644
--- a/etcdserver/config.go
+++ b/etcdserver/config.go
@@ -53,6 +53,7 @@ type ServerConfig struct {
 	BootstrapTimeout time.Duration
 
 	AutoCompactionRetention int
+	AutoCompactionMode      string
 	QuotaBackendBytes       int64
 	MaxTxnOps               uint
 
diff --git a/etcdserver/server.go b/etcdserver/server.go
index ac18b881ec27..100d92bf34ec 100644
--- a/etcdserver/server.go
+++ b/etcdserver/server.go
@@ -221,7 +221,7 @@ type EtcdServer struct {
 
 	SyncTicker *time.Ticker
 	// compactor is used to auto-compact the KV.
-	compactor *compactor.Periodic
+	compactor compactor.Compactor
 
 	// peerRt used to send requests (version, lease) to peers.
 	peerRt   http.RoundTripper
@@ -469,8 +469,11 @@ func NewServer(cfg ServerConfig) (srv *EtcdServer, err error) {
 		return nil, err
 	}
 	srv.authStore = auth.NewAuthStore(srv.be, tp)
-	if h := cfg.AutoCompactionRetention; h != 0 {
-		srv.compactor = compactor.NewPeriodic(h, srv.kv, srv)
+	if num := cfg.AutoCompactionRetention; num != 0 {
+		srv.compactor, err = compactor.New(cfg.AutoCompactionMode, num, srv.kv, srv)
+		if err != nil {
+			return nil, err
+		}
 		srv.compactor.Run()
 	}