From 58c04bc7cb88dcfff6c282570325807ea28b55df Mon Sep 17 00:00:00 2001
From: choury <zhouwei400@gmail.com>
Date: Tue, 9 Apr 2019 20:10:22 +0800
Subject: [PATCH] Support multiple storage backends

---
 cache/memory/memory.go |  8 ++++----
 storagedriver.go       | 21 +++++++++++++--------
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/cache/memory/memory.go b/cache/memory/memory.go
index fbeb2070..99454c77 100644
--- a/cache/memory/memory.go
+++ b/cache/memory/memory.go
@@ -70,7 +70,7 @@ type InMemoryCache struct {
 	lock              sync.RWMutex
 	containerCacheMap map[string]*containerCache
 	maxAge            time.Duration
-	backend           storage.StorageDriver
+	backend           []storage.StorageDriver
 }
 
 func (self *InMemoryCache) AddStats(cInfo *info.ContainerInfo, stats *info.ContainerStats) error {
@@ -86,11 +86,11 @@ func (self *InMemoryCache) AddStats(cInfo *info.ContainerInfo, stats *info.Conta
 		}
 	}()
 
-	if self.backend != nil {
+	for _, backend := range self.backend {
 		// TODO(monnand): To deal with long delay write operations, we
 		// may want to start a pool of goroutines to do write
 		// operations.
-		if err := self.backend.AddStats(cInfo, stats); err != nil {
+		if err := backend.AddStats(cInfo, stats); err != nil {
 			klog.Error(err)
 		}
 	}
@@ -131,7 +131,7 @@ func (self *InMemoryCache) RemoveContainer(containerName string) error {
 
 func New(
 	maxAge time.Duration,
-	backend storage.StorageDriver,
+	backend []storage.StorageDriver,
 ) *InMemoryCache {
 	ret := &InMemoryCache{
 		containerCacheMap: make(map[string]*containerCache, 32),
diff --git a/storagedriver.go b/storagedriver.go
index d59c86a4..60c12261 100644
--- a/storagedriver.go
+++ b/storagedriver.go
@@ -34,19 +34,24 @@ import (
 )
 
 var (
-	storageDriver   = flag.String("storage_driver", "", fmt.Sprintf("Storage `driver` to use. Data is always cached shortly in memory, this controls where data is pushed besides the local cache. Empty means none. Options are: <empty>, %s", strings.Join(storage.ListDrivers(), ", ")))
+	storageDriver   = flag.String("storage_driver", "", fmt.Sprintf("Storage `driver` to use. Data is always cached shortly in memory, this controls where data is pushed besides the local cache. Empty means none, multiple separated by commas. Options are: <empty>, %s", strings.Join(storage.ListDrivers(), ", ")))
 	storageDuration = flag.Duration("storage_duration", 2*time.Minute, "How long to keep data stored (Default: 2min).")
 )
 
 // NewMemoryStorage creates a memory storage with an optional backend storage option.
 func NewMemoryStorage() (*memory.InMemoryCache, error) {
-	backendStorage, err := storage.New(*storageDriver)
-	if err != nil {
-		return nil, err
-	}
-	if *storageDriver != "" {
-		klog.V(1).Infof("Using backend storage type %q", *storageDriver)
+	backendStorages := []storage.StorageDriver{}
+	for _, driver := range strings.Split(*storageDriver, ",") {
+		if driver == "" {
+			continue
+		}
+		storage, err := storage.New(driver)
+		if err != nil {
+			return nil, err
+		}
+		backendStorages = append(backendStorages, storage)
+		klog.V(1).Infof("Using backend storage type %q", driver)
 	}
 	klog.V(1).Infof("Caching stats in memory for %v", *storageDuration)
-	return memory.New(*storageDuration, backendStorage), nil
+	return memory.New(*storageDuration, backendStorages), nil
 }