/*
 * Copyright 2022 CloudWeGo 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 nphttp2

import (
	"net"
	"sync/atomic"
	"testing"
	"time"

	"github.com/bytedance/sonic"

	"github.com/cloudwego/kitex/internal/test"
	"github.com/cloudwego/kitex/pkg/remote"
)

func TestConnPool(t *testing.T) {
	// mock init
	connPool := newMockConnPool()
	defer connPool.Close()
	opt := newMockConnOption()
	ctx := newMockCtxWithRPCInfo()

	// test Get()
	_, err := connPool.Get(ctx, "tcp", mockAddr0, opt)
	test.Assert(t, err == nil, err)

	// test connection reuse
	// keep create new connection until filling the pool size,
	// then Get() will reuse established connection by round-robin
	for i := 0; uint32(i) < connPool.size*2; i++ {
		_, err := connPool.Get(ctx, "tcp", mockAddr1, opt)
		test.Assert(t, err == nil, err)
	}

	// test TimeoutGet()
	opt.ConnectTimeout = time.Microsecond
	_, err = connPool.Get(ctx, "tcp", mockAddr0, opt)
	test.Assert(t, err != nil, err)

	// test Dump()
	dump := connPool.Dump()
	test.Assert(t, dump != nil)
	m := dump.(map[string]interface{})
	test.Assert(t, m[mockAddr0] != nil)
	test.Assert(t, m[mockAddr1] != nil)
	_, err = sonic.Marshal(dump)
	test.Assert(t, err == nil, err)

	// test Clean()
	connPool.Clean("tcp", mockAddr0)
	_, ok := connPool.conns.Load(mockAddr0)
	test.Assert(t, !ok)
}

func TestReleaseConn(t *testing.T) {
	// short connection
	// the connection will be released when put back to the pool
	var closed1 int32 = 0
	dialFunc1 := func(network, address string, timeout time.Duration) (net.Conn, error) {
		npConn := newMockNpConn(address)
		npConn.CloseFunc = func() error {
			atomic.StoreInt32(&closed1, 1)
			return nil
		}
		npConn.mockSettingFrame()
		return npConn, nil
	}
	shortCP := newMockConnPool()
	shortCP.connOpts.ShortConn = true
	defer shortCP.Close()

	conn, err := shortCP.Get(newMockCtxWithRPCInfo(), "tcp", mockAddr0, remote.ConnOption{Dialer: newMockDialerWithDialFunc(dialFunc1)})
	// close stream to ensure no active stream on this connection,
	// which will be released when put back to the connection pool and closed by GracefulClose
	s := conn.(*clientConn).s
	conn.(*clientConn).tr.CloseStream(s, nil)
	test.Assert(t, err == nil, err)
	time.Sleep(100 * time.Millisecond)
	shortCP.Put(conn)
	test.Assert(t, atomic.LoadInt32(&closed1) == 1)

	// long connection
	// the connection will not be released when put back to the pool
	var closed2 int32 = 0
	dialFunc2 := func(network, address string, timeout time.Duration) (net.Conn, error) {
		npConn := newMockNpConn(address)
		npConn.CloseFunc = func() error {
			atomic.StoreInt32(&closed2, 1)
			return nil
		}
		npConn.mockSettingFrame()
		return npConn, nil
	}
	longCP := newMockConnPool()
	longCP.connOpts.ShortConn = false
	defer longCP.Close()

	longConn, err := longCP.Get(newMockCtxWithRPCInfo(), "tcp", mockAddr0, remote.ConnOption{Dialer: newMockDialerWithDialFunc(dialFunc2)})
	longConn.Close()
	test.Assert(t, err == nil, err)
	longCP.Put(longConn)
	test.Assert(t, atomic.LoadInt32(&closed2) == 0)
}
