diff --git a/internal/bbsim/devices/olt.go b/internal/bbsim/devices/olt.go
index 746a58b..305a7ce 100644
--- a/internal/bbsim/devices/olt.go
+++ b/internal/bbsim/devices/olt.go
@@ -309,9 +309,6 @@
 		return err
 	}
 
-	time.Sleep(1 * time.Second) // we need to give the OLT the time to respond to all the pending gRPC request before stopping the server
-	o.StopOltServer()
-
 	if softReboot {
 		for _, pon := range o.Pons {
 			if pon.InternalState.Current() == "enabled" {
@@ -341,11 +338,13 @@
 		}
 	}
 
+	time.Sleep(1 * time.Second) // we need to give the OLT the time to respond to all the pending gRPC request before stopping the server
+	o.StopOltServer()
+
 	// terminate the OLT's processOltMessages go routine
 	close(o.channel)
 
 	o.enableContextCancel()
-
 	time.Sleep(time.Duration(rebootDelay) * time.Second)
 
 	if err := o.InternalState.Event("initialize"); err != nil {
diff --git a/internal/bbsim/devices/onu.go b/internal/bbsim/devices/onu.go
index 549b663..1afb624 100644
--- a/internal/bbsim/devices/onu.go
+++ b/internal/bbsim/devices/onu.go
@@ -342,6 +342,12 @@
 		"stream":  stream,
 	}).Debug("Starting ONU Indication Channel")
 
+	defer onuLogger.WithFields(log.Fields{
+		"onuID":  o.ID,
+		"onuSN":  o.Sn(),
+		"stream": stream,
+	}).Debug("Stopped handling ONU Indication Channel")
+
 loop:
 	for {
 		select {
@@ -493,11 +499,6 @@
 			}
 		}
 	}
-	onuLogger.WithFields(log.Fields{
-		"onuID":  o.ID,
-		"onuSN":  o.Sn(),
-		"stream": stream,
-	}).Debug("Stopped handling ONU Indication Channel")
 }
 
 func NewSN(oltid int, intfid uint32, onuid uint32) *openolt.SerialNumber {