ndm sds

handlers/scraper/remote.go

@@ -4,7 +4,7 @@ import (
 	"errors"
 	"fmt"
 	"net"
-	"syscall"
+	"sync"
 	"time"
 
 	"github.com/kelindar/binary"
@@ -44,44 +44,44 @@ func handleConnection(conn net.Conn) {
 	if err != nil {
 		return
 	}
+	defer session.Close()
 
+	var wg sync.WaitGroup
 	for {
 		stream, err := session.AcceptStream()
 		if err != nil {
 			break
 		}
 
-		go func() {
-			for {
-				select {
-				case rm := <-inChan:
-					buf := []byte(rm.instaData.PostID)
-					if _, err = stream.Write(buf); err != nil {
-						rm.outChan <- err
-						if errors.Is(err, syscall.EPIPE) {
-							return
-						}
-						continue
-					}
-
-					outBuf := make([]byte, 1024*1024)
-					if n, err := stream.Read(outBuf); err == nil {
-						err = binary.Unmarshal(outBuf[:n], rm.instaData)
-						rm.outChan <- err
-					} else {
-						rm.outChan <- err
-						if errors.Is(err, syscall.EPIPE) {
-							return
-						}
-						continue
-					}
-
-				case <-stream.GetDieCh():
+		wg.Add(1)
+		go func(stream *smux.Stream) {
+			defer wg.Done()
+			defer stream.Close()
+
+			for rm := range inChan {
+				buf := []byte(rm.instaData.PostID)
+				if _, err = stream.Write(buf); err != nil {
+					rm.outChan <- err
+					return
+				}
+
+				outBuf := make([]byte, 1024*1024) // Smaller, more reasonable buffer
+				n, err := stream.Read(outBuf)
+				if err != nil {
+					rm.outChan <- err
 					return
 				}
+
+				if err = binary.Unmarshal(outBuf[:n], rm.instaData); err != nil {
+					rm.outChan <- err
+					return
+				}
+				rm.outChan <- nil
 			}
-		}()
+		}(stream)
 	}
+
+	wg.Wait()
 }
 
 func ScrapeRemote(i *InstaData) error {