github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/cmd/redo/apply.go (about)

     1  // Copyright 2021 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 redo
    15  
    16  import (
    17  	"net/http"
    18  	_ "net/http/pprof" // init pprof
    19  	"net/url"
    20  	"runtime/debug"
    21  	"time"
    22  
    23  	"github.com/pingcap/log"
    24  	"github.com/pingcap/tiflow/pkg/applier"
    25  	cmdcontext "github.com/pingcap/tiflow/pkg/cmd/context"
    26  	cerror "github.com/pingcap/tiflow/pkg/errors"
    27  	"github.com/pingcap/tiflow/pkg/util"
    28  	"github.com/spf13/cobra"
    29  	"go.uber.org/zap"
    30  )
    31  
    32  // applyRedoOptions defines flags for the `redo apply` command.
    33  type applyRedoOptions struct {
    34  	options
    35  	sinkURI              string
    36  	enableProfiling      bool
    37  	memoryLimitInGiBytes int64
    38  }
    39  
    40  // newapplyRedoOptions creates new applyRedoOptions for the `redo apply` command.
    41  func newapplyRedoOptions() *applyRedoOptions {
    42  	return &applyRedoOptions{}
    43  }
    44  
    45  // addFlags receives a *cobra.Command reference and binds
    46  // flags related to template printing to it.
    47  func (o *applyRedoOptions) addFlags(cmd *cobra.Command) {
    48  	cmd.Flags().StringVar(&o.sinkURI, "sink-uri", "", "target database sink-uri")
    49  	// the possible error returned from MarkFlagRequired is `no such flag`
    50  	cmd.MarkFlagRequired("sink-uri") //nolint:errcheck
    51  	cmd.Flags().BoolVar(&o.enableProfiling, "enable-profiling", true, "enable pprof profiling")
    52  	cmd.Flags().Int64Var(&o.memoryLimitInGiBytes, "memory-limit", 10, "memory limit in GiB")
    53  }
    54  
    55  //nolint:unparam
    56  func (o *applyRedoOptions) complete(cmd *cobra.Command) error {
    57  	// parse sinkURI as a URI
    58  	sinkURI, err := url.Parse(o.sinkURI)
    59  	if err != nil {
    60  		return cerror.WrapError(cerror.ErrSinkURIInvalid, err)
    61  	}
    62  	rawQuery := sinkURI.Query()
    63  	// set safe-mode to true if not set
    64  	if rawQuery.Get("safe-mode") != "true" {
    65  		rawQuery.Set("safe-mode", "true")
    66  		sinkURI.RawQuery = rawQuery.Encode()
    67  		o.sinkURI = sinkURI.String()
    68  	}
    69  
    70  	totalMemory, err := util.GetMemoryLimit()
    71  	if err == nil {
    72  		totalMemoryInBytes := int64(float64(totalMemory) * 0.8)
    73  		memoryLimitInBytes := o.memoryLimitInGiBytes * 1024 * 1024 * 1024
    74  		if totalMemoryInBytes != 0 && memoryLimitInBytes > totalMemoryInBytes {
    75  			memoryLimitInBytes = totalMemoryInBytes
    76  		}
    77  		debug.SetMemoryLimit(memoryLimitInBytes)
    78  		log.Info("set memory limit", zap.Int64("memoryLimit", memoryLimitInBytes))
    79  	}
    80  
    81  	return nil
    82  }
    83  
    84  // run runs the `redo apply` command.
    85  func (o *applyRedoOptions) run(cmd *cobra.Command) error {
    86  	ctx := cmdcontext.GetDefaultContext()
    87  
    88  	if o.enableProfiling {
    89  		go func() {
    90  			server := &http.Server{
    91  				Addr:              ":6060",
    92  				ReadHeaderTimeout: 5 * time.Second,
    93  			}
    94  			log.Info("Start http pprof server", zap.String("addr", server.Addr))
    95  			if err := server.ListenAndServe(); err != nil {
    96  				log.Fatal("http pprof", zap.Error(err))
    97  			}
    98  		}()
    99  	}
   100  
   101  	cfg := &applier.RedoApplierConfig{
   102  		Storage: o.storage,
   103  		SinkURI: o.sinkURI,
   104  		Dir:     o.dir,
   105  	}
   106  	ap := applier.NewRedoApplier(cfg)
   107  	err := ap.Apply(ctx)
   108  	if err != nil {
   109  		return err
   110  	}
   111  	cmd.Println("Apply redo log successfully")
   112  	return nil
   113  }
   114  
   115  // newCmdApply creates the `redo apply` command.
   116  func newCmdApply(opt *options) *cobra.Command {
   117  	o := newapplyRedoOptions()
   118  	command := &cobra.Command{
   119  		Use:   "apply",
   120  		Short: "Apply redo logs in target sink",
   121  		RunE: func(cmd *cobra.Command, args []string) error {
   122  			o.options = *opt
   123  			if err := o.complete(cmd); err != nil {
   124  				return err
   125  			}
   126  			return o.run(cmd)
   127  		},
   128  	}
   129  	o.addFlags(command)
   130  
   131  	return command
   132  }