package mppkill_test

import (
	"context"
	"database/sql"
	"time"

	. "github.com/onsi/ginkgo"
	"github.com/pingcap/log"
	"github.com/pingcap/test-infra/sdk/resource"
	"golang.org/x/sync/errgroup"

	"github.com/pingcap/endless/pkg/client/tidb"
	"github.com/pingcap/endless/pkg/matcher"
	"github.com/pingcap/endless/pkg/matcher/metrics"
	"github.com/pingcap/endless/pkg/util"
	"github.com/pingcap/endless/testcase/tiflash/pkg"
	"github.com/pingcap/endless/testcase/tiflash/pkg/workload/ch"
)

/*
多并发执行 mpp 查询，定时 kill 这些标记为 killmpp 的查询
单并发执行 mpp 查询，该查询不会被 kill，期望不会报错
*/

var _ = Describe("TiFlash", func() {
	var tc resource.TiDBCluster
	var chWorkload *ch.ChWorkload
	BeforeEach(func() {
		tc = suiteTestCtx.Resource("tc").(resource.TiDBCluster)
		chWorkload = ch.NewWorkload(suiteTestCtx).(*ch.ChWorkload)
		chWorkload.Prepare(suiteTestCtx)
	})

	Context("cancel mpp query", func() {
		It("should not crash when kill connection #001#", func() {
			// stable run 30 minutes before start kill mpp
			ctx0, _ := context.WithTimeout(context.Background(), 30.*time.Minute)
			g0, ctxr0 := errgroup.WithContext(ctx0)
			for currency := 0; currency < 3; currency++ {
				g0.Go(util.WithGinkgoRecover(func() {
					matcher.Try(chWorkload.RunAPWithComment(ctxr0, pkg.CtxDoneHandler, ""))
				}))
			}
			matcher.Try(g0.Wait())

			ctxr, _ := context.WithTimeout(context.Background(), 5.*time.Hour)
			db := matcher.Try(pkg.OpenDB(tc, chWorkload.DBName(), pkg.OlapConParams...)).(*sql.DB)
			tidbClient := tidb.NewClient(db)
			g, ctxr2 := errgroup.WithContext(ctxr)
			g.Go(util.WithGinkgoRecover(func() {
				chWorkload.RunTP(ctxr2, func(error) error { return nil })
			}))
			for currency := 0; currency < 2; currency++ {
				g.Go(util.WithGinkgoRecover(func() {
					// Error 1317: Query execution was interrupted
					chWorkload.RunAPWithComment(ctxr2, func(error) error { return nil }, "/*+ killmpp */")
				}))
			}
			g.Go(util.WithGinkgoRecover(func() {
				matcher.Try(chWorkload.RunAPWithComment(ctxr2, pkg.ErrorHandler(pkg.CtxDoneHandler, pkg.TiFlashPrimaryLockHandler), "/*+ stablerun */"))
			}))
			g.Go(util.WithGinkgoRecover(func() {
				for {
					if pkg.Sleep(ctxr2, 10*time.Second) {
						return
					}
					log.Info("start kill connection id")
					if err := tidbClient.KillProcesslistByKeyword(ctxr2, "query", "killmpp"); err != nil {
						log.Info(err.Error())
					}
				}
			}))
			g.Go(util.WithGinkgoRecover(func() {
				matcher.ConsistentlyWithMetrics(tc).
					ShouldNot(metrics.TiFlashRestart().WithCtx(ctxr2))
			}))
			matcher.Try(g.Wait())
		})
		// Deprecated: same with high currency case2, so deprecate it
		It("should not crash when out of memory useage #002#", func() {
			ctxr, _ := context.WithTimeout(context.Background(), 5.*time.Hour)
			g, ctxr2 := errgroup.WithContext(ctxr)
			// Exchange receiver meet error : DB::TiFlashException: Memory limit (total) exceeded
			for currency := 0; currency < 2; currency++ {
				g.Go(util.WithGinkgoRecover(func() {
					matcher.Try(chWorkload.RunAP(ctxr2, func(error) error { return nil }))
				}))
			}
			g.Go(util.WithGinkgoRecover(func() {
				chWorkload.RunTP(ctxr2, func(error) error { return nil })
			}))
			g.Go(util.WithGinkgoRecover(func() {
				matcher.ConsistentlyWithMetrics(tc).
					ShouldNot(metrics.TiFlashRestart().WithCtx(ctxr2))
			}))
			matcher.Try(g.Wait())
		})
	})
})