/*
 *
 * Copyright 2018 gRPC authors.
 *
 * 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 grpclb

import (
	"fmt"
	"sync"
	"testing"
	"time"

	"google.golang.org/grpc/balancer"
	"google.golang.org/grpc/resolver"
)

type mockSubConn struct {
	balancer.SubConn
}

type mockClientConn struct {
	balancer.ClientConn

	mu       sync.Mutex
	subConns map[balancer.SubConn]resolver.Address
}

func newMockClientConn() *mockClientConn {
	return &mockClientConn{
		subConns: make(map[balancer.SubConn]resolver.Address),
	}
}

func (mcc *mockClientConn) NewSubConn(addrs []resolver.Address, opts balancer.NewSubConnOptions) (balancer.SubConn, error) {
	sc := &mockSubConn{}
	mcc.mu.Lock()
	defer mcc.mu.Unlock()
	mcc.subConns[sc] = addrs[0]
	return sc, nil
}

func (mcc *mockClientConn) RemoveSubConn(sc balancer.SubConn) {
	mcc.mu.Lock()
	defer mcc.mu.Unlock()
	delete(mcc.subConns, sc)
}

const testCacheTimeout = 100 * time.Millisecond

func checkMockCC(mcc *mockClientConn, scLen int) error {
	mcc.mu.Lock()
	defer mcc.mu.Unlock()
	if len(mcc.subConns) != scLen {
		return fmt.Errorf("mcc = %+v, want len(mcc.subConns) = %v", mcc.subConns, scLen)
	}
	return nil
}

func checkCacheCC(ccc *lbCacheClientConn, sccLen, sctaLen int) error {
	ccc.mu.Lock()
	defer ccc.mu.Unlock()
	if len(ccc.subConnCache) != sccLen {
		return fmt.Errorf("ccc = %+v, want len(ccc.subConnCache) = %v", ccc.subConnCache, sccLen)
	}
	if len(ccc.subConnToAddr) != sctaLen {
		return fmt.Errorf("ccc = %+v, want len(ccc.subConnToAddr) = %v", ccc.subConnToAddr, sctaLen)
	}
	return nil
}

// Test that SubConn won't be immediately removed.
func TestLBCacheClientConnExpire(t *testing.T) {
	mcc := newMockClientConn()
	if err := checkMockCC(mcc, 0); err != nil {
		t.Fatal(err)
	}

	ccc := newLBCacheClientConn(mcc)
	ccc.timeout = testCacheTimeout
	if err := checkCacheCC(ccc, 0, 0); err != nil {
		t.Fatal(err)
	}

	sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
	// One subconn in MockCC.
	if err := checkMockCC(mcc, 1); err != nil {
		t.Fatal(err)
	}
	// No subconn being deleted, and one in CacheCC.
	if err := checkCacheCC(ccc, 0, 1); err != nil {
		t.Fatal(err)
	}

	ccc.RemoveSubConn(sc)
	// One subconn in MockCC before timeout.
	if err := checkMockCC(mcc, 1); err != nil {
		t.Fatal(err)
	}
	// One subconn being deleted, and one in CacheCC.
	if err := checkCacheCC(ccc, 1, 1); err != nil {
		t.Fatal(err)
	}

	// Should all become empty after timeout.
	var err error
	for i := 0; i < 2; i++ {
		time.Sleep(testCacheTimeout)
		err = checkMockCC(mcc, 0)
		if err != nil {
			continue
		}
		err = checkCacheCC(ccc, 0, 0)
		if err != nil {
			continue
		}
	}
	if err != nil {
		t.Fatal(err)
	}
}

// Test that NewSubConn with the same address of a SubConn being removed will
// reuse the SubConn and cancel the removing.
func TestLBCacheClientConnReuse(t *testing.T) {
	mcc := newMockClientConn()
	if err := checkMockCC(mcc, 0); err != nil {
		t.Fatal(err)
	}

	ccc := newLBCacheClientConn(mcc)
	ccc.timeout = testCacheTimeout
	if err := checkCacheCC(ccc, 0, 0); err != nil {
		t.Fatal(err)
	}

	sc, _ := ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
	// One subconn in MockCC.
	if err := checkMockCC(mcc, 1); err != nil {
		t.Fatal(err)
	}
	// No subconn being deleted, and one in CacheCC.
	if err := checkCacheCC(ccc, 0, 1); err != nil {
		t.Fatal(err)
	}

	ccc.RemoveSubConn(sc)
	// One subconn in MockCC before timeout.
	if err := checkMockCC(mcc, 1); err != nil {
		t.Fatal(err)
	}
	// One subconn being deleted, and one in CacheCC.
	if err := checkCacheCC(ccc, 1, 1); err != nil {
		t.Fatal(err)
	}

	// Recreate the old subconn, this should cancel the deleting process.
	sc, _ = ccc.NewSubConn([]resolver.Address{{Addr: "address1"}}, balancer.NewSubConnOptions{})
	// One subconn in MockCC.
	if err := checkMockCC(mcc, 1); err != nil {
		t.Fatal(err)
	}
	// No subconn being deleted, and one in CacheCC.
	if err := checkCacheCC(ccc, 0, 1); err != nil {
		t.Fatal(err)
	}

	var err error
	// Should not become empty after 2*timeout.
	time.Sleep(2 * testCacheTimeout)
	err = checkMockCC(mcc, 1)
	if err != nil {
		t.Fatal(err)
	}
	err = checkCacheCC(ccc, 0, 1)
	if err != nil {
		t.Fatal(err)
	}

	// Call remove again, will delete after timeout.
	ccc.RemoveSubConn(sc)
	// One subconn in MockCC before timeout.
	if err := checkMockCC(mcc, 1); err != nil {
		t.Fatal(err)
	}
	// One subconn being deleted, and one in CacheCC.
	if err := checkCacheCC(ccc, 1, 1); err != nil {
		t.Fatal(err)
	}

	// Should all become empty after timeout.
	for i := 0; i < 2; i++ {
		time.Sleep(testCacheTimeout)
		err = checkMockCC(mcc, 0)
		if err != nil {
			continue
		}
		err = checkCacheCC(ccc, 0, 0)
		if err != nil {
			continue
		}
	}
	if err != nil {
		t.Fatal(err)
	}
}