stop waiting on context cancel

cmd/windows_agent/service.go

@@ -77,8 +77,9 @@ func (p *agentService) runLoop() {
 		return
 	}
 
+	go p.startPing()
+
 	for {
-		go p.startPing()
 		p.run()
 		wgDone := utils.WaitChan(&p.wg)
 
@@ -146,14 +147,15 @@ waitUrl:
 		}
 
 		p.wg.Add(1)
-
-		go sftp.Serve(p.ctx, &p.wg, sftpConfig, "0.0.0.0", port, driveLetter)
+		go func() {
+			sftp.Serve(p.ctx, sftpConfig, "0.0.0.0", port, driveLetter)
+			p.wg.Done()
+		}()
 	}
 }
 
 func (p *agentService) Stop(s service.Service) error {
 	p.cancel()
 
-	p.wg.Wait()
 	return nil
 }

internal/agent/sftp/sftp.go

@@ -9,16 +9,13 @@ import (
 	"net"
 	"net/url"
 	"strings"
-	"sync"
 
 	"github.com/pkg/sftp"
 	"github.com/sonroyaalmerol/pbs-plus/internal/agent/snapshots"
 	"golang.org/x/crypto/ssh"
 )
 
-func Serve(ctx context.Context, wg *sync.WaitGroup, sftpConfig *SFTPConfig, address, port string, driveLetter string) {
-	defer wg.Done()
-
+func Serve(ctx context.Context, sftpConfig *SFTPConfig, address, port string, driveLetter string) {
 	listenAt := fmt.Sprintf("%s:%s", address, port)
 	listener, err := net.Listen("tcp", listenAt)
 	if err != nil {