Add keep alive for soap.RoundTripper

session/keep_alive.go

@@ -0,0 +1,114 @@
+/*
+Copyright (c) 2014 VMware, Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package session
+
+import (
+	"sync"
+	"time"
+
+	"github.com/vmware/govmomi/vim25/methods"
+	"github.com/vmware/govmomi/vim25/soap"
+	"golang.org/x/net/context"
+)
+
+type keepAlive struct {
+	sync.Mutex
+
+	roundTripper  soap.RoundTripper
+	idleTime      time.Duration
+	notifyRequest chan struct{}
+	notifyStop    chan struct{}
+
+	// keepAlive executes a request in the background with the purpose of
+	// keeping the session active. The response for this request is discarded.
+	keepAlive func(soap.RoundTripper)
+}
+
+func defaultKeepAlive(roundTripper soap.RoundTripper) {
+	methods.GetServiceContent(context.Background(), roundTripper)
+}
+
+// KeepAlive wraps the specified soap.RoundTripper and executes a meaningless
+// API request in the background after the RoundTripper has been idle for the
+// specified amount of idle time. The keep alive process only starts once a
+// user logs in and runs until the user logs out again.
+func KeepAlive(roundTripper soap.RoundTripper, idleTime time.Duration) soap.RoundTripper {
+	k := &keepAlive{
+		roundTripper:  roundTripper,
+		idleTime:      idleTime,
+		notifyRequest: make(chan struct{}),
+	}
+
+	k.keepAlive = defaultKeepAlive
+
+	return k
+}
+
+func (k *keepAlive) start() {
+	k.Lock()
+	defer k.Unlock()
+
+	if k.notifyStop != nil {
+		return
+	}
+
+	// This channel must be closed to terminate idle timer.
+	k.notifyStop = make(chan struct{})
+
+	go func() {
+		t := time.NewTimer(k.idleTime)
+
+		for {
+			select {
+			case <-k.notifyStop:
+				return
+			case <-k.notifyRequest:
+				t.Reset(k.idleTime)
+			case <-t.C:
+				k.keepAlive(k.roundTripper)
+				t = time.NewTimer(k.idleTime)
+			}
+		}
+	}()
+}
+
+func (k *keepAlive) stop() {
+	k.Lock()
+	defer k.Unlock()
+
+	if k.notifyStop != nil {
+		close(k.notifyStop)
+		k.notifyStop = nil
+	}
+}
+
+func (k *keepAlive) RoundTrip(ctx context.Context, req, res soap.HasFault) error {
+	err := k.roundTripper.RoundTrip(ctx, req, res)
+	if err != nil {
+		return err
+	}
+
+	// Start ticker on login, stop ticker on logout.
+	switch req.(type) {
+	case *methods.LoginBody:
+		k.start()
+	case *methods.LogoutBody:
+		k.stop()
+	}
+
+	return nil
+}

session/keep_alive_test.go

@@ -0,0 +1,154 @@
+/*
+Copyright (c) 2014 VMware, Inc. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package session
+
+import (
+	"fmt"
+	"net/url"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/vmware/govmomi/test"
+	"github.com/vmware/govmomi/vim25/methods"
+	"github.com/vmware/govmomi/vim25/soap"
+	"golang.org/x/net/context"
+)
+
+type testKeepAlive int
+
+func (t *testKeepAlive) Func(soap.RoundTripper) {
+	*t++
+}
+
+func newManager(t *testing.T) (*Manager, *url.URL) {
+	u := test.URL()
+	if u == nil {
+		t.SkipNow()
+	}
+
+	soapClient := soap.NewClient(u, true)
+	serviceContent, err := methods.GetServiceContent(context.Background(), soapClient)
+	if err != nil {
+		t.Error(err)
+	}
+
+	return NewManager(soapClient, serviceContent), u
+}
+
+func TestKeepAlive(t *testing.T) {
+	var i testKeepAlive
+	var j int
+
+	m, u := newManager(t)
+	k := KeepAlive(m.roundTripper, time.Millisecond)
+	k.(*keepAlive).keepAlive = i.Func
+	m.roundTripper = k
+
+	// Expect keep alive to not have triggered yet
+	if i != 0 {
+		t.Errorf("Expected i == 0, got i: %d", i)
+	}
+
+	// Logging in starts keep alive
+	err := m.Login(context.Background(), u.User)
+	if err != nil {
+		t.Error(err)
+	}
+
+	time.Sleep(2 * time.Millisecond)
+
+	// Expect keep alive to triggered at least once
+	if i == 0 {
+		t.Errorf("Expected i != 0, got i: %d", i)
+	}
+
+	j = int(i)
+	time.Sleep(2 * time.Millisecond)
+
+	// Expect keep alive to triggered at least once more
+	if int(i) <= j {
+		t.Errorf("Expected i > j, got i: %d, j: %d", i, j)
+	}
+
+	// Logging out stops keep alive
+	err = m.Logout(context.Background())
+	if err != nil {
+		t.Error(err)
+	}
+
+	j = int(i)
+	time.Sleep(2 * time.Millisecond)
+
+	// Expect keep alive to have stopped
+	if int(i) != j {
+		t.Errorf("Expected i == j, got i: %d, j: %d", i, j)
+	}
+}
+
+func testSessionOK(t *testing.T, m *Manager, ok bool) {
+	s, err := m.UserSession(context.Background())
+	if err != nil {
+		t.Error(err)
+	}
+
+	if ok && s == nil {
+		t.Error("Expected session to be OK, but is invalid")
+	}
+
+	if !ok && s != nil {
+		t.Error("Expected session to be invalid, but is OK")
+	}
+}
+
+// Run with:
+//
+//   env GOVMOMI_KEEPALIVE_TEST=1 go test -timeout=60m -run TestRealKeepAlive
+//
+func TestRealKeepAlive(t *testing.T) {
+	if os.Getenv("GOVMOMI_KEEPALIVE_TEST") != "1" {
+		t.SkipNow()
+	}
+
+	m1, u1 := newManager(t)
+	m2, u2 := newManager(t)
+
+	// Enable keepalive on m2
+	k := KeepAlive(m2.roundTripper, 10*time.Minute)
+	m2.roundTripper = k
+
+	// Logging in starts keep alive
+	if err := m1.Login(context.Background(), u1.User); err != nil {
+		t.Error(err)
+	}
+	if err := m2.Login(context.Background(), u2.User); err != nil {
+		t.Error(err)
+	}
+
+	// Expect both sessions to be non-nil
+	testSessionOK(t, m1, true)
+	testSessionOK(t, m2, true)
+
+	// Wait for m1 to time out
+	delay := 31 * time.Minute
+	fmt.Printf("%s: Waiting %d minutes for session to time out...\n", time.Now(), int(delay.Minutes()))
+	time.Sleep(delay)
+
+	// Expect m1's session to be invalid, m2's session to be valid
+	testSessionOK(t, m1, false)
+	testSessionOK(t, m2, true)
+}