CatapultClient.Close() fails consistently due to race condition

[avanaur]

Whenever I tried Close(), it always fails with the error below.

panic: close tcp 192.168.20.11:64395->52.221.198.188:3000: use of closed network connection

goroutine 69 [running]:
github.com/proximax-storage/go-xpx-catapult-sdk/sdk/websocket.(*CatapultWebsocketClientImpl).Listen(0xc000388000)
	C:/Users/Tyrone/go/src/github.com/proximax-storage/go-xpx-catapult-sdk/sdk/websocket/client.go:162 +0x24a
created by main.testWS
	C:/Users/Tyrone/go/src/github.com/proximax-storage/csd-test-playground/test_websocket/main.go:56 +0x1a0

My understanding is there is a race condition on Close() and Listen(). Upon Close(), Listen() would immediately run the first case below.

	for {
		select {
		case <-c.ctx.Done():
			if c.conn != nil {
				if err := c.conn.Close(); err != nil {
					panic(err)
				}
				c.conn = nil
			}
			return
		case msg := <-messagesChan:
			go c.messageRouter.RouteMessage(msg)
		}
	}

My code:

	conf, err := sdk.NewConfig([]string{"http://csddev1.xpxsirius.io:3000", "http://csddev2.xpxsirius.io:3000"}, sdk.PrivateTest, sdk.WebsocketReconnectionDefaultTimeout)
	if err != nil {
		panic(err)
	}

	ctx, cancel := context.WithTimeout(context.TODO(), time.Second*30)
	wsClient, err := websocket.NewClient(ctx, conf)
	if err != nil {
		panic(err)
	}
	go wsClient.Listen()

	address, err := sdk.NewAddressFromRaw("WAFSWQIF2UBLBHOFOYMIQJXZKRSTNONZIM6VYAZT")
	if err != nil {
		panic(err)
	}

	txn, err := sdk.NewTransferTransaction(
		sdk.NewDeadline(time.Hour),
		address,
		[]*sdk.Mosaic{},
		sdk.NewPlainMessage(fmt.Sprintf("hi there %d", rand.Int())),
		sdk.PrivateTest,
	)
	if err != nil {
		panic(err)
	}

	acct, err := sdk.NewAccountFromPrivateKey("b007b006c2178a8819c6b38dd3b49db1c2170483fdfcff21656f071c7541b5a9", sdk.PrivateTest)
	if err != nil {
		panic(err)
	}

	signedTxn, err := acct.Sign(txn)
	if err != nil {
		panic(err)
	}

	var wg sync.WaitGroup
	wg.Add(1)
	err = wsClient.AddUnconfirmedAddedHandlers(address, func(txn sdk.Transaction) bool {
		fmt.Println("FOUND", txn.GetAbstractTransaction().Hash.String())
		if txn.GetAbstractTransaction().Hash.String() == signedTxn.Hash.String() {
			fmt.Println("FOUND HASH IT ON UNCONFIRMED", txn.GetAbstractTransaction().Hash.String())
			wg.Done()
		}
		return false
	})
	if err != nil {
		panic(err)
	}

	client := sdk.NewClient(nil, conf)
	_, err = client.Transaction.Announce(context.Background(), signedTxn)
	fmt.Println("ANNOUNCED", signedTxn.Hash)
	if err != nil {
		panic(err)
	}

	wg.Wait()

	wsClient.Close()

My workaround right now is to use context's cancel function. It works most of the time but has a separate issue.