From b73d03776163466e6ed9a11d08670d78f408dd46 Mon Sep 17 00:00:00 2001
From: James Bardin <j.bardin@gmail.com>
Date: Thu, 25 May 2017 09:18:42 -0400
Subject: [PATCH] have StateHook periodically PersistState

Have StateHook periodically call PersistState to flush any cached state
to permanent storage. This uses a minimal 10 second interval between
calls to PersistState.
---
 backend/local/hook_state.go      | 26 +++++++++
 backend/local/hook_state_test.go | 97 ++++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+)

diff --git a/backend/local/hook_state.go b/backend/local/hook_state.go
index 5483c4344..62e33416e 100644
--- a/backend/local/hook_state.go
+++ b/backend/local/hook_state.go
@@ -2,17 +2,27 @@ package local
 
 import (
 	"sync"
+	"time"
 
 	"github.com/hashicorp/terraform/state"
 	"github.com/hashicorp/terraform/terraform"
 )
 
+// interval between forced PersistState calls by StateHook
+const persistStateHookInterval = 10 * time.Second
+
 // StateHook is a hook that continuously updates the state by calling
 // WriteState on a state.State.
 type StateHook struct {
 	terraform.NilHook
 	sync.Mutex
 
+	// lastPersist is the time of the last call to PersistState, for periodic
+	// updates to remote state. PostStateUpdate will force a call PersistState
+	// if it has been more that persistStateHookInterval since the last call to
+	// PersistState.
+	lastPersist time.Time
+
 	State state.State
 }
 
@@ -26,8 +36,24 @@ func (h *StateHook) PostStateUpdate(
 		if err := h.State.WriteState(s); err != nil {
 			return terraform.HookActionHalt, err
 		}
+
+		// periodically persist the state
+		if time.Since(h.lastPersist) > persistStateHookInterval {
+			if err := h.persistState(); err != nil {
+				return terraform.HookActionHalt, err
+			}
+		}
 	}
 
 	// Continue forth
 	return terraform.HookActionContinue, nil
 }
+
+func (h *StateHook) persistState() error {
+	if h.State != nil {
+		err := h.State.PersistState()
+		h.lastPersist = time.Now()
+		return err
+	}
+	return nil
+}
diff --git a/backend/local/hook_state_test.go b/backend/local/hook_state_test.go
index 7f4d88770..1614942d6 100644
--- a/backend/local/hook_state_test.go
+++ b/backend/local/hook_state_test.go
@@ -1,7 +1,9 @@
 package local
 
 import (
+	"sync"
 	"testing"
+	"time"
 
 	"github.com/hashicorp/terraform/state"
 	"github.com/hashicorp/terraform/terraform"
@@ -27,3 +29,98 @@ func TestStateHook(t *testing.T) {
 		t.Fatalf("bad state: %#v", is.State())
 	}
 }
+
+// testPersistState stores the state on WriteState, and
+type testPersistState struct {
+	*state.InmemState
+
+	mu        sync.Mutex
+	persisted bool
+}
+
+func (s *testPersistState) WriteState(state *terraform.State) error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.persisted = false
+	return s.InmemState.WriteState(state)
+}
+
+func (s *testPersistState) PersistState() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	s.persisted = true
+	return nil
+}
+
+// verify that StateHook calls PersistState if the last call was more than
+// persistStateHookInterval
+func TestStateHookPersist(t *testing.T) {
+	is := &testPersistState{
+		InmemState: &state.InmemState{},
+	}
+	hook := &StateHook{State: is}
+
+	s := state.TestStateInitial()
+	hook.PostStateUpdate(s)
+
+	// the first call should persist, since the last time was zero
+	if !is.persisted {
+		t.Fatal("PersistState not called")
+	}
+
+	s.Serial++
+	hook.PostStateUpdate(s)
+
+	// this call should not have persisted
+	if is.persisted {
+		t.Fatal("PostStateUpdate called PersistState early")
+	}
+
+	if !is.State().Equal(s) {
+		t.Fatalf("bad state: %#v", is.State())
+	}
+
+	// set the last call back to before our interval
+	hook.lastPersist = time.Now().Add(-2 * persistStateHookInterval)
+
+	s.Serial++
+	hook.PostStateUpdate(s)
+
+	if !is.persisted {
+		t.Fatal("PersistState not called")
+	}
+
+	if !is.State().Equal(s) {
+		t.Fatalf("bad state: %#v", is.State())
+	}
+}
+
+// verify that the satet hook is safe for concurrent use
+func TestStateHookRace(t *testing.T) {
+	is := &state.InmemState{}
+	var hook terraform.Hook = &StateHook{State: is}
+
+	s := state.TestStateInitial()
+
+	var wg sync.WaitGroup
+
+	for i := 0; i < 100; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			action, err := hook.PostStateUpdate(s)
+			if err != nil {
+				t.Fatalf("err: %s", err)
+			}
+			if action != terraform.HookActionContinue {
+				t.Fatalf("bad: %v", action)
+			}
+			if !is.State().Equal(s) {
+				t.Fatalf("bad state: %#v", is.State())
+			}
+		}()
+	}
+	wg.Wait()
+}