im_server/server/cluster_order_test.go

package main

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

	"github.com/tinode/chat/server/concurrency"
	"github.com/tinode/chat/server/store/types"
)

type orderedTopicProxy struct {
	mu       sync.Mutex
	received []int64
	done     chan struct{}
}

func (proxy *orderedTopicProxy) TopicProxy(msg *ClusterResp, rejected *bool) error {
	if msg.SrvMsg.Data.SeqId == 1 {
		time.Sleep(50 * time.Millisecond)
	}
	proxy.mu.Lock()
	proxy.received = append(proxy.received, msg.SrvMsg.Data.SeqId)
	if len(proxy.received) == 2 {
		close(proxy.done)
	}
	proxy.mu.Unlock()
	return nil
}

func TestClusterWriteLoopPreservesTopicOrder(t *testing.T) {
	serverConn, clientConn := net.Pipe()
	rpcServer := rpc.NewServer()
	proxy := &orderedTopicProxy{done: make(chan struct{})}
	if err := rpcServer.RegisterName("Cluster", proxy); err != nil {
		t.Fatal(err)
	}
	go rpcServer.ServeConn(serverConn)

	pool := concurrency.NewGoRoutinePool(4)
	previousCluster := globals.cluster
	globals.cluster = &Cluster{proxyEventQueue: pool}
	t.Cleanup(func() {
		globals.cluster = previousCluster
		_ = clientConn.Close()
		_ = serverConn.Close()
		pool.Stop()
	})

	sess := &Session{
		sid:          "multiplex-order-test",
		proto:        MULTIPLEX,
		proxiedTopic: "p2p-order-test",
		send:         make(chan any, 2),
		clnode:       &ClusterNode{endpoint: rpc.NewClient(clientConn), connected: true},
	}
	tenantID := types.TenantID(1)
	sess.send <- &ServerComMessage{TenantID: tenantID, Data: &MsgServerData{SeqId: 1}}
	sess.send <- &ServerComMessage{TenantID: tenantID, Data: &MsgServerData{SeqId: 2}}

	// Scheduling twice used to start two concurrent drains for the same topic.
	sess.scheduleClusterWriteLoop()
	sess.scheduleClusterWriteLoop()

	select {
	case <-proxy.done:
	case <-time.After(2 * time.Second):
		t.Fatal("timed out waiting for cluster responses")
	}
	deadline := time.Now().Add(2 * time.Second)
	for sess.clusterWriteScheduled.Load() && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if sess.clusterWriteScheduled.Load() {
		t.Fatal("cluster write task did not release session ownership")
	}

	proxy.mu.Lock()
	defer proxy.mu.Unlock()
	if len(proxy.received) != 2 || proxy.received[0] != 1 || proxy.received[1] != 2 {
		t.Fatalf("cluster responses arrived out of order: %v", proxy.received)
	}
}

func TestClusterWriteLoopRejectsWritesAfterTermination(t *testing.T) {
	previousSessionStore := globals.sessionStore
	globals.sessionStore = NewSessionStore(time.Minute)
	t.Cleanup(func() { globals.sessionStore = previousSessionStore })

	sess := &Session{
		sid:    "multiplex-termination-test",
		proto:  MULTIPLEX,
		send:   make(chan any, 1),
		stop:   make(chan any, 1),
		detach: make(chan string, 1),
		subs:   make(map[string]*Subscription),
	}
	sess.stop <- nil

	sess.clusterWriteLoop("grp-termination-test")

	if atomic.LoadInt32(&sess.terminating) == 0 {
		t.Fatal("terminated multiplexing session still accepts writes")
	}
	if !sess.queueOut(&ServerComMessage{TenantID: types.TenantID(1)}) {
		t.Fatal("write to terminated multiplexing session should be discarded")
	}
	if len(sess.send) != 0 {
		t.Fatalf("terminated multiplexing session retained %d outbound messages", len(sess.send))
	}
}

func TestMultiplexEnqueueWaitsForDrain(t *testing.T) {
	sess := &Session{proto: MULTIPLEX, send: make(chan any, 1)}
	sess.send <- "existing"

	go func() {
		time.Sleep(10 * time.Millisecond)
		<-sess.send
	}()

	payload := &ServerComMessage{TenantID: types.TenantID(1)}
	if !sess.enqueueMultiplexOutbound(payload, 200*time.Millisecond) {
		t.Fatal("multiplex enqueue failed while the drain was making progress")
	}
	if got := <-sess.send; got != payload {
		t.Fatalf("unexpected queued payload: %#v", got)
	}
}

func TestMultiplexEnqueueTimesOutWithoutDrain(t *testing.T) {
	sess := &Session{proto: MULTIPLEX, send: make(chan any, 1)}
	sess.send <- "existing"

	started := time.Now()
	if sess.enqueueMultiplexOutbound("next", 20*time.Millisecond) {
		t.Fatal("multiplex enqueue succeeded without queue capacity")
	}
	if elapsed := time.Since(started); elapsed < 15*time.Millisecond || elapsed > 200*time.Millisecond {
		t.Fatalf("multiplex enqueue timeout outside expected bounds: %s", elapsed)
	}
}

func TestMultiplexEnqueueStopsWhenSessionTerminates(t *testing.T) {
	sess := &Session{proto: MULTIPLEX, send: make(chan any, 1)}
	sess.send <- "existing"

	go func() {
		time.Sleep(10 * time.Millisecond)
		sess.clusterQueueMu.Lock()
		atomic.StoreInt32(&sess.terminating, 1)
		<-sess.send
		sess.clusterQueueMu.Unlock()
	}()

	if !sess.enqueueMultiplexOutbound("next", 200*time.Millisecond) {
		t.Fatal("terminating multiplex session should discard the pending write")
	}
	if len(sess.send) != 0 {
		t.Fatalf("terminating multiplex session retained %d outbound messages", len(sess.send))
	}
}