oto/services/core/internal/runnersocket/server_test.go

package runnersocket

import (
	"net"
	"testing"
	"time"

	protoSocket "git.toki-labs.com/toki/proto-socket/go"
	"github.com/toki/oto/services/core/internal/cicdstate"
	"github.com/toki/oto/services/core/internal/runnerregistry"
	otopb "github.com/toki/oto/services/core/oto"
)

func testServer(registry *runnerregistry.Registry) *Server {
	return New("127.0.0.1:0", registry, cicdstate.NewStore())
}

// pipeClient creates a TcpClient backed by one end of a net.Pipe.
// heartbeat is disabled (intervalSec=0) to avoid background timer interference.
// Cleanup closes both ends.
func pipeClient(t *testing.T) (*protoSocket.TcpClient, net.Conn) {
	t.Helper()
	a, b := net.Pipe()
	client := protoSocket.NewTcpClient(a, 0, 0, parserMap())
	t.Cleanup(func() { _ = a.Close(); _ = b.Close() })
	return client, b
}

// TestServerSetClientReplacesDuplicate verifies that registering a second client
// for the same runner ID closes the old one and stores the new one.
func TestServerReplacesDuplicateRunnerClient(t *testing.T) {
	s := testServer(runnerregistry.New())

	client1, _ := pipeClient(t)
	client2, _ := pipeClient(t)

	s.setClient("runner-A", client1)
	s.setClient("runner-A", client2)

	// Allow Close() to propagate through connCloseOnce.
	time.Sleep(20 * time.Millisecond)

	if client1.IsAlive() {
		t.Error("old client must be closed after replacement, but it is still alive")
	}
	if got := s.clientFor("runner-A"); got != client2 {
		t.Errorf("clientFor must return new client after replacement; got %v", got)
	}
}

// TestServerIgnoresStaleClientDisconnect verifies that removeClient ignores
// a client that is not the current entry for the given runner ID.
func TestServerIgnoresStaleClientDisconnect(t *testing.T) {
	s := testServer(runnerregistry.New())

	client1, _ := pipeClient(t)
	client2, _ := pipeClient(t)

	s.setClient("runner-B", client1)

	// client2 is a different object – simulates a stale disconnect listener.
	removed := s.removeClient("runner-B", client2)
	if removed {
		t.Error("removeClient must return false for a client that is not the current one")
	}
	if got := s.clientFor("runner-B"); got != client1 {
		t.Errorf("current client must still be client1 after stale removeClient; got %v", got)
	}
}

// TestServerStaleClientDisconnectDoesNotDisconnectRegistry verifies the complete
// stale-disconnect protection chain:
//
//  1. setClient replaces the old client.
//  2. removeClient returns false for the stale old client.
//  3. Because removeClient returns false, registry.Disconnect is NOT called,
//     so the runner remains in a healthy (non-disconnected) state.
func TestServerStaleClientDisconnectDoesNotDisconnectRegistry(t *testing.T) {
	registry := runnerregistry.New()
	s := testServer(registry)

	resp := registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "tok",
		RunnerId:        "runner-C",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})
	if !resp.GetAccepted() {
		t.Fatalf("registry.Register rejected: %s", resp.GetRejectReason())
	}
	registry.Heartbeat("runner-C", otopb.HeartbeatStatus_HEARTBEAT_STATUS_HEALTHY)

	client1, _ := pipeClient(t)
	client2, _ := pipeClient(t)

	s.setClient("runner-C", client1)
	// Duplicate connection: client2 replaces client1.
	s.setClient("runner-C", client2)

	// Simulate stale client1's disconnect listener.
	removed := s.removeClient("runner-C", client1)
	if removed {
		t.Error("stale removeClient should return false")
	}

	// Because removeClient returned false the disconnect listener guard skips
	// registry.Disconnect. The runner must NOT be in disconnected state.
	runner, ok := registry.Snapshot("runner-C")
	if !ok {
		t.Fatal("runner-C not found in registry")
	}
	if runner.Status == runnerregistry.StatusDisconnected {
		t.Errorf("runner status must not be %q after stale disconnect; got %q",
			runnerregistry.StatusDisconnected, runner.Status)
	}
	if got := s.clientFor("runner-C"); got != client2 {
		t.Errorf("current client must still be client2; got %v", got)
	}
}

// TestServerDuplicateReplaceWithListenerDoesNotDeadlock verifies that replacing
// a client that carries a production-style disconnect listener does not deadlock.
//
// Before the fix, setClient held s.mu while calling old.Close(), which fired
// the listener synchronously. The listener called removeClient → tried to
// acquire s.mu again → deadlock.  After the fix, old.Close() is called after
// s.mu is released, so the re-entry is safe.
func TestServerDuplicateReplaceWithListenerDoesNotDeadlock(t *testing.T) {
	registry := runnerregistry.New()
	s := testServer(registry)

	resp := registry.Register(&otopb.RegisterRunnerRequest{
		EnrollmentToken: "tok",
		RunnerId:        "runner-E",
		ProtocolVersion: "oto.runner.v1",
		Capability:      &otopb.RunnerCapability{Name: "oto-runner", Version: "1.0.0"},
	})
	if !resp.GetAccepted() {
		t.Fatalf("registry.Register rejected: %s", resp.GetRejectReason())
	}

	client1, _ := pipeClient(t)
	client2, _ := pipeClient(t)

	// Mirror the production disconnect listener from attachClientHandlers.
	// This is the re-entry path: removeClient acquires s.mu, which would
	// deadlock if setClient still held s.mu during old.Close().
	client1.AddDisconnectListener(func(c *protoSocket.TcpClient) {
		if s.removeClient("runner-E", c) {
			s.registry.Disconnect("runner-E")
		}
	})

	s.setClient("runner-E", client1)

	// This must complete without hanging. If the lock-order bug is present the
	// goroutine stalls here and the test times out via go test -timeout.
	s.setClient("runner-E", client2)

	// Allow Close() and the disconnect listener to propagate.
	time.Sleep(20 * time.Millisecond)

	if client1.IsAlive() {
		t.Error("old client must be closed after replacement")
	}
	if got := s.clientFor("runner-E"); got != client2 {
		t.Errorf("clientFor must return new client after replacement; got %v", got)
	}
	// The stale listener must not have disconnected the runner.
	runner, ok := registry.Snapshot("runner-E")
	if !ok {
		t.Fatal("runner-E not found in registry after replacement")
	}
	if runner.Status == runnerregistry.StatusDisconnected {
		t.Errorf("runner must not be disconnected after stale listener fires; got %q", runner.Status)
	}
}

// TestServerHeartbeatUsesRegisteredRunnerID verifies that connectedRunnerIDs
// reflects only currently active clients and that clientFor returns nil after
// the current client is removed.
func TestServerHeartbeatUsesRegisteredRunnerID(t *testing.T) {
	s := testServer(runnerregistry.New())

	client1, _ := pipeClient(t)

	if ids := s.connectedRunnerIDs(); len(ids) != 0 {
		t.Errorf("expected no connected runners before any setClient; got %v", ids)
	}

	s.setClient("runner-D", client1)

	ids := s.connectedRunnerIDs()
	if len(ids) != 1 || ids[0] != "runner-D" {
		t.Errorf("expected [runner-D] after setClient; got %v", ids)
	}

	// Current client removed → runner no longer in connected list.
	s.removeClient("runner-D", client1)

	if ids := s.connectedRunnerIDs(); len(ids) != 0 {
		t.Errorf("expected no connected runners after removeClient; got %v", ids)
	}
	if got := s.clientFor("runner-D"); got != nil {
		t.Errorf("expected nil from clientFor after removeClient; got %v", got)
	}
}