The test relies on time.Sleep() calls to wait for message processing to complete (lines 155, 205, 264, 376). This makes tests non-deterministic and unnecessarily slow. Consider using synchronization primitives like channels or WaitGroups in the mock handler to signal when processing is complete, rather than arbitrary sleep durations.

Change: Refactor scheduler_test.go to use synchronization primitives instead of time.Sleep

1. Eliminated time.Sleep: Removed all arbitrary sleep calls (time.Sleep) used to wait for async processing.
2. Implemented sync.WaitGroup: Introduced sync.WaitGroup in the test functions and mock handlers to track the exact completion of task processing.

The Worker's wg WaitGroup is added to when workers start but never waited upon during shutdown. Each worker calls wg.Done() when it exits, but there's no corresponding wg.Wait() in the shutdown sequence. This means the Shutdown() method can return before all workers have finished processing their remaining tasks and exited cleanly. The UEScheduler should store a reference to this WaitGroup and wait for all workers to complete during shutdown.

Change: Integrated sync.WaitGroup into the UEScheduler struct and ensured Shutdown() waits for all workers. Removed the unused stopChan and Stop() method.
Rationale: This ensures the AMF process does not exit until all workers have finished processing their remaining tasks.

// internal/ngap/scheduler.go

type UEScheduler struct {
    // ...
    wg sync.WaitGroup // <--- Stored in struct
}

func NewUEScheduler(...) *UEScheduler {
    // ...
    for i := 0; i < numWorkers; i++ {
        // Pass the scheduler's WG reference to workers
        scheduler.workers[i] = NewWorker(i, taskBufferSize, handler, &scheduler.wg)
    }
    return scheduler
}

func (s *UEScheduler) Shutdown() {
    s.workerMutex.Lock()
    // ... close channels ...
    s.workerMutex.Unlock()

    s.wg.Wait() // <--- Explicit wait prevents data loss
    logger.NgapLog.Info("All workers shut down successfully")
}

I’d like to double-check whether the current implementation can run into a send on closed channel issue.

Shutdown() closes w.taskChan, while UEScheduler.DispatchTask() may still be running in a different goroutine and calling worker.Submit(task). If these overlap, it seems possible for a send to occur after the channel has been closed.

Could you help confirm whether this race is already prevented by the current design, or if additional coordination (e.g. a scheduler-level stop signal or submit-side guarding) is needed?

After consideration, I have decided to revert to the stopChan + select design.

Implementation Details:

• Mechanism: taskChan is never explicitly closed (left to GC). Submit uses select to manage both shutdown behavior and normal traffic flow, ensuring backpressure is maintained.
• Draining: A drainAndExit function is added to process residual packets during shutdown.
• Safety: This approach resolves both the "send on closed channel" panic and "shutdown deadlock" issues.
• Lock-Free: All mutexes have been removed. The worker structure is effectively read-only after initialization, and the shutdown process no longer depends on the specific closing order of taskChan.
• Error Handling: defer recover is retained in the run loop strictly to handle unexpected runtime errors (e.g., nil pointer dereferences) and prevent the worker from crashing.

Performance Impact (Hot Path Analysis):
I acknowledge that using select in the Submit hot path introduces overhead. I conducted a simple test using UERANSIM with 100 concurrent UEs (-n 100) to measure the impact on registration time (calculated via log timestamps):

• The initial PR version (which had no extra overhead in Submit): 10.292 seconds
• Current Version (Select-based): 11.829 seconds

Result: The current safety-focused design incurs an approximate 14% performance cost.

The Submit method performs a blocking send on the task channel. If the worker's buffer is full and no workers are consuming (e.g., during shutdown or if a worker panics), this will block indefinitely, potentially deadlocking the system. Consider using a select statement with a timeout or a default case to handle the scenario where the buffer is full, allowing the caller to handle backpressure appropriately.

Change: Moved the blocking Submit call outside the workerMutex lock scope.
Rationale: This prevents a deadlock scenario where a full channel blocks Submit (holding the RLock), preventing Shutdown from acquiring the Lock. Note that I retained the blocking behavior (instead of dropping tasks) to maintain backpressure for flow control in the 5G control plane.

// internal/ngap/scheduler.go

func (s *UEScheduler) DispatchTask(task Task) {
    s.workerMutex.RLock()
    // Hash the UE ID to determine which worker should handle it
    workerIndex := s.hashUEID(task.UEID)
    worker := s.workers[workerIndex]
    s.workerMutex.RUnlock() // <--- Unlock BEFORE submitting

    logger.NgapLog.Debugf("Dispatching UE ID %d to Worker %d", task.UEID, workerIndex)
    worker.Submit(task)     // <--- Safe blocking submission
}

My understanding is that changing only DispatchTask does not fully address the issue. If w.taskChan is full, UEScheduler can still be blocked on worker.Submit(task).

I would suggest that the select-based handling inside Submit is still necessary, so that cases like a full buffer or stopChan being closed are explicitly handled. This also ensures proper behavior during worker shutdown.

I have re-implemented the select-based logic within Submit. Currently, when the worker is full, the operation blocks to ensure backpressure, which effectively reduces retransmissions and prevents signaling storms. If a shutdown signal is received while the worker is full, Submit will exit immediately, leaving drainAndExit to handle the remaining packets.

When the scheduler is not initialized, the code falls back to sequential processing with a warning log. However, there's no mechanism to prevent continued operation in this degraded state. If scheduler initialization fails during startup, the system silently continues without parallelization. Consider making scheduler initialization a hard requirement, or at least log this as an error rather than a warning to make the degraded state more visible.

The byte conversion for hashing uses little-endian order (shifting right and taking lower bytes first). While this is not incorrect, using a standard library function like binary.LittleEndian.PutUint64() would be clearer and more maintainable. This makes the byte order explicit and follows Go conventions.

Change: Replaced fnv hashing with simple modulo arithmetic.
Rationale: fnv required memory allocation (make([]byte)) on every packet. Using modulo is zero-allocation and sufficient for load distribution.

// internal/ngap/scheduler.go

func (s *UEScheduler) hashUEID(ueID uint64) int {
    return int(ueID % uint64(s.numWorkers))
}

The stopChan is closed in the Stop() method but never actually used in the Shutdown() function. The Shutdown method only closes the taskChan, but doesn't call Stop() on individual workers. This means the stopChan case in the worker's run() loop is dead code. Either remove the stopChan or properly utilize it in the shutdown sequence.

#194 (comment)
I removed the unused stopChan and Stop() method. Worker termination is now handled purely by closing the taskChan, which is a cleaner and more idiomatic Go pattern.

The configuration getters return 0 for auto-detect when NgapWorkerPoolSize is not set or is 0. However, the InitScheduler function is called with this 0 value, and within InitScheduler, it checks if numWorkers <= 0 and defaults to runtime.NumCPU(). This logic is duplicated - InitScheduler already handles the auto-detect case, so the comment on line 1043 is misleading as the function doesn't directly perform auto-detection based on the returned 0.

Updated the comment in config.go to clearly explain that a return value of 0 signals the caller to perform auto-detection.

If InitScheduler returns an error (from the initErr variable), the code calls logger.InitLog.Fatalf which terminates the program. However, InitScheduler never actually sets initErr to a non-nil value - the function always returns nil. Either remove the error return type from InitScheduler or properly implement error handling for initialization failures (e.g., if worker creation fails).

Change: Removed the error return from InitScheduler.
Rationale: Since invalid inputs (<=0) are now handled by applying safe defaults (auto-detect or 4096 buffer), the function cannot fail. Removing the return value clarifies the API and satisfies linter checks.

// internal/ngap/scheduler.go

func InitScheduler(numWorkers int, taskBufferSize int, handler func(conn net.Conn, msg []byte)) {
    globalSchedulerOnce.Do(func() {
        if numWorkers <= 0 { numWorkers = runtime.NumCPU() }
        if taskBufferSize <= 0 { taskBufferSize = 4096 }
        // ... init ...
    })
    // No return error needed
}

The comment "Import the ngap package to access the scheduler" is misleading. The import statement is at the top of the file (line 12), not at this location. This comment should be removed or clarified to describe what this function does, not how imports work.

Removed the confusing comment in service.go

The hash function creates a new FNV hash instance and allocates a new byte slice on every call. Since this is called for every message dispatch (potentially thousands of times per second), this creates unnecessary allocations and garbage collection pressure. Consider creating a pool of hash instances or using a simpler modulo operation directly on the uint64 value if perfect distribution isn't critical: return int(ueID % uint64(s.numWorkers)).

#194 (comment)
Switched to modulo arithmetic to avoid memory allocation on every packet dispatch.

Should this be SourceAMFUENGAPID? Otherwise, it will use a worker different from that of existing RanUe.

Yeah, this possible issue I have written in above Notes. Under the current global AMF Context architecture, a UE message triggers a worker switch only once. While this results in a cache miss, it does not lead to any functional errors. However, if the global Context is to be removed in the future, the allocation mechanisms for both amf-ue-id and ran-ue-id must be modified accordingly.

Architecture & Safety Considerations:
To minimize changes to the core architecture, this implementation continues to utilize the global AMFContext (via sync.Map) for storing UE contexts. Consequently, a UE's InitialUEMessage (keyed by RAN ID) may be processed by Worker A, while subsequent NAS messages (keyed by AMF ID) may be processed by Worker B.

• Safety: This approach is safe because the global context uses thread-safe sync.Map for access. Furthermore, the 5G Request-Response architecture prevents race conditions where a UE sends messages with a RAN ID and an AMF ID simultaneously.
• Future Work: Achieving a strictly local "Per-UE Connection" model (removing global lock contention entirely) would require removing the global UePool and modifying the AMF-UE-ID allocation mechanism to bind specific IDs to specific workers.

My understanding (please correct me if I’m mistaken) is that while InitialUEMessage only triggers a single worker switch and is functionally safe under the current global AMFContext + sync.Map design, PathSwitchRequest is different.

At this stage, the UE may already have ongoing traffic (e.g. previous uplink NGAP messages) keyed by AMF UE NGAP ID. If PathSwitchRequest is dispatched using the new RAN UE NGAP ID instead of SourceAMFUENGAPID, it could introduce an additional worker switch during an active UE lifecycle.

A possible scenario is:

Initial UE Message   → Worker A (RAN ID)
Uplink NAS Transport → Worker B (AMF UE ID)
Path Switch Request → Worker C (new RAN UE ID)

This could lead to multiple workers accessing the same UE context and extra worker hopping.
If my understanding differs from the actual implementation, I’m happy to discuss and clarify.

Sorry, I didn't fully grasp your question at first. You are absolutely right—we should keep the same Source AMF UE NGAP ID. I've updated the code. Thanks!