broke up session loop

pool.go

@@ -119,9 +119,9 @@ func WithEventsBufferSize(value int) PoolOption {
 	}
 }
 
-func WithSessionConfig(wc SessionConfig) PoolOption {
+func WithSessionConfig(cfg SessionConfig) PoolOption {
 	return func(c *PoolConfig) {
-		c.SessionConfig = wc
+		c.SessionConfig = cfg
 	}
 }
 
@@ -356,8 +356,8 @@ func (p *Pool) Connect(id string, opts ...ConnectOption) error {
 		dialFn = o.dialFn
 	}
 
-	wc := p.config.SessionConfig
+	sc := p.config.SessionConfig
-	session, err := newSession(p.ctx, id, dialFn, &wc, p.handler)
+	session, err := newSession(p.ctx, id, dialFn, &sc, p.handler)
 	if err != nil {
 		return err
 	}

session.go

@@ -181,7 +181,18 @@ func (w *session) Start(pool PoolPlugin) {
 	}
 
 	var wg sync.WaitGroup
-	wg.Go(func() { w.run(w.ctx, pool) })
+	wg.Go(func() {
+		for {
+			conn, ok := w.dial(w.ctx, pool)
+			if !ok {
+				return
+			}
+			w.serve(w.ctx, conn, pool)
+			if !w.reset(w.ctx) {
+				return
+			}
+		}
+	})
 
 	if w.logger != nil {
 		w.logger.Debug("started")
@@ -194,8 +205,9 @@ func (w *session) Start(pool PoolPlugin) {
 	}
 }
 
-func (w *session) run(ctx context.Context, pool PoolPlugin) {
+func (w *session) dial(
-	for {
+	ctx context.Context, pool PoolPlugin,
+) (*transport.Connection, bool) {
 	var retryMgr *transport.RetryManager
 	if !w.config.RetryDisabled {
 		retryMgr = transport.NewRetryManager(w.config.RetryConfig)
@@ -217,19 +229,25 @@ func (w *session) run(ctx context.Context, pool PoolPlugin) {
 			!errors.Is(err, context.DeadlineExceeded) {
 			pool.Retire(err)
 		}
-			return
+		return nil, false
 	}
 
-		// setup new connection
 	conn, _ := transport.NewConnection(
 		ctx, socket, &w.config.ConnectionConfig, w.handler)
-		w.conn.Store(conn)
+	return conn, true
-		pool.Events <- PoolEvent{ID: w.url, Kind: EventConnected, At: time.Now()}
+}
 
+func (w *session) serve(
+	ctx context.Context, conn *transport.Connection, pool PoolPlugin,
+) {
 	if w.logger != nil {
 		w.logger.Debug("session: started")
 	}
 
+	// setup connection
+	w.conn.Store(conn)
+	pool.Events <- PoolEvent{ID: w.url, Kind: EventConnected, At: time.Now()}
+
 	// start keepalive service
 	stopKeepalive, inactive, heartbeat := w.startKeepalive()
 
@@ -280,28 +298,28 @@ func (w *session) run(ctx context.Context, pool PoolPlugin) {
 	conn.Close()
 	stopKeepalive()
 
-		if w.logger != nil {
-			w.logger.Info("disconnected")
-		}
-		if w.logger != nil {
-			w.logger.Debug("session: ended")
-		}
-
 	// tear down connection
 	w.conn.Store(nil)
 	pool.Events <- PoolEvent{ID: w.url, Kind: EventDisconnected, At: time.Now()}
 
+	if w.logger != nil {
+		w.logger.Info("disconnected")
+		w.logger.Debug("session: ended")
+	}
+}
+
+func (w *session) reset(ctx context.Context) bool {
 	// exit if session is shutting down
 	select {
 	case <-ctx.Done():
-			return
+		return false
 	default:
 	}
 
 	// refresh session
 	time.Sleep(w.config.ReconnectDelay)
 	w.restartCount.Add(1)
-	}
+	return true
 }
 
 func (w *session) startKeepalive() (