From 24fe84808302a8dbb1d1cae0af4806d456aee1d6 Mon Sep 17 00:00:00 2001
From: Louis Royer <infos.louis.royer@gmail.com>
Date: Thu, 12 Dec 2024 11:31:29 +0100
Subject: [PATCH] Add `WaitReady(ctx) error` to PFCPEntityInterface

---
 README.md                    | 10 ++++++----
 pfcp/api/entity_interface.go |  2 ++
 pfcp/entity.go               | 15 +++++++++++++++
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/README.md b/README.md
index 37ad774..a7a482a 100644
--- a/README.md
+++ b/README.md
@@ -10,9 +10,10 @@
 ### UPF
 
 ```golang
+ctx := context.TODO()
 upNode := NewPFCPEntityUP(UPF_NODE_ID, UPF_IP_ADDR) // node id can be an IP Address or a FQDN
-go upNode.ListenAndServe()
-defer upNode.Close()
+go upNode.ListenAndServeContext(ctx)
+upnode.WaitReady(ctx)
 // Access list of associations
 associations := upNode.GetPFCPAssociations()
 // Access list of sessions
@@ -22,9 +23,10 @@ sessions := upNode.GetPFCPSessions()
 ### SMF
 
 ```golang
+ctx := context.TODO()
 cpNode := NewPFCPEntityCP(SMF_NODE_ID, SMF_IP_ADDR) // node id can be an IP Address or a FQDN
-go cpNode.ListenAndServe()
-defer cpNode.Close()
+go cpNode.ListenAndServeContext(ctx)
+cpNode.WaitReady(ctx)
 association, _ := cpNode.NewEstablishedPFCPAssociation(ie.NewNodeIDHeuristic(UPFADDR))
 session, _ := a.CreateSession(pdrs, fars)
 
diff --git a/pfcp/api/entity_interface.go b/pfcp/api/entity_interface.go
index 6fd8eef..59bed94 100644
--- a/pfcp/api/entity_interface.go
+++ b/pfcp/api/entity_interface.go
@@ -6,6 +6,7 @@
 package api
 
 import (
+	"context"
 	"net/netip"
 
 	"github.com/wmnsk/go-pfcp/ie"
@@ -25,4 +26,5 @@ type PFCPEntityInterface interface {
 	AddEstablishedPFCPSession(session PFCPSessionInterface) error
 	LogPFCPRules()
 	Options() EntityOptionsInterface
+	WaitReady(ctx context.Context) error
 }
diff --git a/pfcp/entity.go b/pfcp/entity.go
index f1b3144..b0cffe5 100644
--- a/pfcp/entity.go
+++ b/pfcp/entity.go
@@ -38,6 +38,9 @@ type PFCPEntity struct {
 	mu        sync.Mutex
 	pfcpConns []*onceClosePfcpConn
 	closeFunc context.CancelFunc
+
+	// check start
+	waitReady chan struct{}
 }
 
 // onceClosePfcpConn wraps a PfcpConn, protecting it from multiple Close calls.
@@ -133,6 +136,7 @@ func NewPFCPEntity(nodeID string, listenAddr netip.Addr, kind string, handlers m
 		kind:              kind,
 		options:           options,
 		pfcpConns:         nil,
+		waitReady:         make(chan struct{}),
 	}
 }
 
@@ -224,6 +228,16 @@ func (e *PFCPEntity) ListenAndServeContext(ctx context.Context) error {
 	}
 }
 
+// Blocks until the PFCPEntity is ready or context is done
+func (e *PFCPEntity) WaitReady(ctx context.Context) error {
+	select {
+	case <-e.waitReady:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
 // Run the entity with the provided context.
 // Always return a non-nil error and close the PFCPConn.
 func (e *PFCPEntity) Serve(ctx context.Context, conn *PFCPConn) error {
@@ -240,6 +254,7 @@ func (e *PFCPEntity) Serve(ctx context.Context, conn *PFCPConn) error {
 		select {
 		case <-serveCtx.Done():
 			// Stop signal received
+			close(e.waitReady)
 			return serveCtx.Err()
 		default:
 			buf := make([]byte, pfcputil.DEFAULT_MTU) // TODO: get MTU of interface instead of using DEFAULT_MTU