github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/tests/mq_protocol_tests/framework/dsl.go (about)

     1  // Copyright 2020 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package framework
    15  
    16  import (
    17  	"context"
    18  	"time"
    19  
    20  	"github.com/cenkalti/backoff/v4"
    21  	"github.com/pingcap/errors"
    22  	"github.com/pingcap/log"
    23  	"go.uber.org/zap"
    24  )
    25  
    26  const (
    27  	waitMaxPollInterval = time.Second * 5
    28  )
    29  
    30  // Awaitable represents the handle of an SQL operation that can be waited on
    31  type Awaitable interface {
    32  	SetTimeOut(duration time.Duration) Awaitable
    33  	Wait() Checkable
    34  }
    35  
    36  // Checkable represents the handle of an SQL operation whose correctness can be checked
    37  type Checkable interface {
    38  	Check() error
    39  }
    40  
    41  type pollable interface {
    42  	poll(ctx context.Context) (bool, error)
    43  }
    44  
    45  type pollableAndCheckable interface {
    46  	pollable
    47  	Checkable
    48  }
    49  
    50  type errorCheckableAndAwaitable struct {
    51  	error
    52  }
    53  
    54  // Check implements Checkable
    55  func (e *errorCheckableAndAwaitable) Check() error {
    56  	return e.error
    57  }
    58  
    59  // Wait implements Awaitable
    60  func (e *errorCheckableAndAwaitable) Wait() Checkable {
    61  	return e
    62  }
    63  
    64  // SetTimeOut implements Awaitable
    65  func (e *errorCheckableAndAwaitable) SetTimeOut(duration time.Duration) Awaitable {
    66  	return e
    67  }
    68  
    69  type basicAwaitable struct {
    70  	pollableAndCheckable
    71  	timeout time.Duration
    72  }
    73  
    74  // SetTimeOut implements Awaitable
    75  func (b *basicAwaitable) SetTimeOut(duration time.Duration) Awaitable {
    76  	b.timeout = duration
    77  	return b
    78  }
    79  
    80  // Wait implements Awaitable
    81  func (b *basicAwaitable) Wait() Checkable {
    82  	var (
    83  		ctx    context.Context
    84  		cancel context.CancelFunc
    85  	)
    86  	if b.timeout == 0 {
    87  		ctx, cancel = context.WithCancel(context.Background())
    88  	} else {
    89  		ctx, cancel = context.WithTimeout(context.Background(), b.timeout)
    90  	}
    91  	defer cancel()
    92  
    93  	expBackoff := backoff.NewExponentialBackOff()
    94  	expBackoff.MaxInterval = waitMaxPollInterval
    95  	for {
    96  		select {
    97  		case <-ctx.Done():
    98  			return &errorCheckableAndAwaitable{ctx.Err()}
    99  		default:
   100  		}
   101  
   102  		ok, err := b.poll(ctx)
   103  		if err != nil {
   104  			return &errorCheckableAndAwaitable{errors.Annotate(err, "Wait() failed with error")}
   105  		}
   106  
   107  		if ok {
   108  			log.Debug("Wait(): pollable finished")
   109  			return b
   110  		}
   111  
   112  		interval := expBackoff.NextBackOff()
   113  		if interval == backoff.Stop {
   114  			return &errorCheckableAndAwaitable{errors.New("Maximum retry interval reached")}
   115  		}
   116  		log.Debug("Wait(): pollable returned false, backing off", zap.Duration("interval", interval))
   117  
   118  		ch := time.After(interval)
   119  		select {
   120  		case <-ctx.Done():
   121  			return &errorCheckableAndAwaitable{ctx.Err()}
   122  		case <-ch:
   123  		}
   124  	}
   125  }