owpengram-server/internal/mtprotoedge/rpc_delivery_hook_executor_test.go

package mtprotoedge

import (
	"errors"
	"sync/atomic"
	"testing"
	"time"
)

func hookTestMessage(reqMsgID int64, coordinator *rpcResultDeliveryCoordinator, fn func()) *encodedOutboundMessage {
	msg := &encodedOutboundMessage{delivery: newRPCResultDelivery(reqMsgID, coordinator)}
	msg.setDeliveryHook(fn)
	return msg
}

func TestRPCDeliveryHookExecutorBoundsAdmissionWithoutBlockingDelivery(t *testing.T) {
	executor := newRPCDeliveryHookExecutor(1, 1)
	started := make(chan struct{})
	release := make(chan struct{})
	first := hookTestMessage(1, nil, func() {
		close(started)
		<-release
	})
	if err := first.prepareDeliveryHook(executor); err != nil {
		t.Fatalf("reserve first hook: %v", err)
	}
	delivered := make(chan struct{})
	go func() {
		first.markDelivered()
		close(delivered)
	}()
	select {
	case <-delivered:
	case <-time.After(100 * time.Millisecond):
		t.Fatal("physical delivery blocked on hook execution")
	}
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("first hook did not start")
	}

	second := hookTestMessage(2, nil, func() {})
	if err := second.prepareDeliveryHook(executor); !errors.Is(err, ErrRPCDeliveryHookCapacity) {
		t.Fatalf("second reservation = %v, want capacity error", err)
	}
	close(release)
	deadline := time.Now().Add(time.Second)
	for len(executor.slots) != 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if got := len(executor.slots); got != 0 {
		t.Fatalf("executor retained %d capacity slots", got)
	}
	snapshot := executor.runtimeSnapshot()
	if snapshot.workers != 1 || snapshot.capacity != 1 || snapshot.completed != 1 || snapshot.rejected != 1 ||
		snapshot.reserved != 0 || snapshot.queued != 0 || snapshot.running != 0 || snapshot.durationSeconds <= 0 {
		t.Fatalf("executor snapshot = %#v", snapshot)
	}
}

func TestRPCDeliveryHookExecutorIsolatesPanicsAndContinues(t *testing.T) {
	executor := newRPCDeliveryHookExecutor(1, 2)
	first := hookTestMessage(1, nil, func() { panic("hook boom") })
	done := make(chan struct{})
	second := hookTestMessage(2, nil, func() { close(done) })
	if err := first.prepareDeliveryHook(executor); err != nil {
		t.Fatalf("reserve panic hook: %v", err)
	}
	if err := second.prepareDeliveryHook(executor); err != nil {
		t.Fatalf("reserve following hook: %v", err)
	}
	first.markDelivered()
	second.markDelivered()
	select {
	case <-done:
	case <-time.After(time.Second):
		t.Fatal("worker stopped after a hook panic")
	}
	deadline := time.Now().Add(time.Second)
	for executor.panics.Load() != 1 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if got := executor.panics.Load(); got != 1 {
		t.Fatalf("recorded hook panics = %d, want 1", got)
	}
}

func TestEquivalentRPCDeliveryAttemptsShareExactlyOnceCoordinator(t *testing.T) {
	executor := newRPCDeliveryHookExecutor(1, 2)
	var calls atomic.Int32
	first := hookTestMessage(11, nil, func() { calls.Add(1) })
	second := hookTestMessage(22, first.delivery.coordinator, nil)
	if err := first.prepareDeliveryHook(executor); err != nil {
		t.Fatalf("reserve first attempt: %v", err)
	}
	if err := second.prepareDeliveryHook(executor); err != nil {
		t.Fatalf("reserve equivalent attempt: %v", err)
	}
	first.markDelivered()
	second.markDelivered()
	deadline := time.Now().Add(time.Second)
	for calls.Load() != 1 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if got := calls.Load(); got != 1 {
		t.Fatalf("equivalent delivery hooks = %d, want 1", got)
	}
	deadline = time.Now().Add(time.Second)
	for len(executor.slots) != 0 && time.Now().Before(deadline) {
		time.Sleep(time.Millisecond)
	}
	if got := len(executor.slots); got != 0 {
		t.Fatalf("equivalent attempts leaked %d tickets", got)
	}
}

func TestRPCDeliveryHookExecutorStopRejectsNewAndDrainsReserved(t *testing.T) {
	executor := newRPCDeliveryHookExecutor(1, 2)
	ticket, ok := executor.reserve()
	if !ok {
		t.Fatal("reserve ticket")
	}
	stopped := make(chan bool, 1)
	go func() { stopped <- executor.stop(time.Second) }()
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		executor.mu.Lock()
		stopping := executor.stopping
		executor.mu.Unlock()
		if stopping {
			break
		}
		time.Sleep(time.Millisecond)
	}
	if _, ok := executor.reserve(); ok {
		t.Fatal("executor accepted a reservation after stop")
	}
	select {
	case <-stopped:
		t.Fatal("stop returned while a reserved ticket was still owned")
	default:
	}
	ticket.release()
	select {
	case ok := <-stopped:
		if !ok {
			t.Fatal("executor did not drain before timeout")
		}
	case <-time.After(time.Second):
		t.Fatal("executor stop did not finish after ticket release")
	}
	snapshot := executor.runtimeSnapshot()
	if snapshot.reserved != 0 || snapshot.queued != 0 || snapshot.running != 0 || snapshot.rejected < 1 {
		t.Fatalf("stopped executor snapshot = %#v", snapshot)
	}
}

func TestRPCDeliveryHookExecutorIsServerScoped(t *testing.T) {
	first := New(Options{RPCDeliveryHookWorkers: 2, RPCDeliveryHookMaxPending: 7})
	second := New(Options{RPCDeliveryHookWorkers: 3, RPCDeliveryHookMaxPending: 9})
	if first.rpcDeliveryHooks == nil || second.rpcDeliveryHooks == nil || first.rpcDeliveryHooks == second.rpcDeliveryHooks {
		t.Fatal("servers did not receive isolated delivery-hook executors")
	}
	firstSnapshot := first.RuntimeSnapshot()
	secondSnapshot := second.RuntimeSnapshot()
	if firstSnapshot.RPCDeliveryHookWorkers != 2 || firstSnapshot.RPCDeliveryHookCapacity != 7 ||
		secondSnapshot.RPCDeliveryHookWorkers != 3 || secondSnapshot.RPCDeliveryHookCapacity != 9 {
		t.Fatalf("server delivery-hook limits = %#v / %#v", firstSnapshot, secondSnapshot)
	}
}