github.com/alibaba/ilogtail/pkg@v0.0.0-20250526110833-c53b480d046c/helper/dumper.go (about)

     1  // Copyright 2023 iLogtail Authors
     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  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package helper
    16  
    17  import (
    18  	"bytes"
    19  	"context"
    20  	"encoding/binary"
    21  	"encoding/json"
    22  	"io"
    23  	"os"
    24  	"path"
    25  	"sync"
    26  	"time"
    27  
    28  	"github.com/alibaba/ilogtail/pkg/config"
    29  	"github.com/alibaba/ilogtail/pkg/helper/async"
    30  	"github.com/alibaba/ilogtail/pkg/logger"
    31  )
    32  
    33  type DumpDataReq struct {
    34  	Body   []byte
    35  	URL    string
    36  	Header map[string][]string
    37  }
    38  type DumpDataResp struct {
    39  	Body   []byte
    40  	Header map[string]string
    41  }
    42  
    43  // DumpData current only for http protocol
    44  type DumpData struct {
    45  	Req  DumpDataReq
    46  	Resp DumpDataResp
    47  }
    48  
    49  type Dumper struct {
    50  	input             chan *DumpData
    51  	dumpDataKeepFiles []string
    52  	prefix            string
    53  	maxKeepFiles      int
    54  	wg                sync.WaitGroup
    55  	stop              chan struct{}
    56  	// unittest
    57  	writeCounter *async.UnitTestCounterHelper
    58  }
    59  
    60  func (d *Dumper) Init() {
    61  	// 只有 service_http_server 插件会使用这个模块
    62  	_ = os.MkdirAll(path.Join(config.LoongcollectorGlobalConfig.LoongCollectorDebugDir, "dump"), 0750)
    63  	d.input = make(chan *DumpData, 10)
    64  	d.stop = make(chan struct{})
    65  	files, err := GetFileListByPrefix(path.Join(config.LoongcollectorGlobalConfig.LoongCollectorDebugDir, "dump"), d.prefix, true, 0)
    66  	if err != nil {
    67  		logger.Warning(context.Background(), "LIST_HISTORY_DUMP_ALARM", "err", err)
    68  	} else {
    69  		d.dumpDataKeepFiles = files
    70  	}
    71  
    72  }
    73  
    74  func (d *Dumper) Start() {
    75  	if d == nil {
    76  		return
    77  	}
    78  	d.wg.Add(1)
    79  	go func() {
    80  		defer d.wg.Done()
    81  		d.doDumpFile()
    82  	}()
    83  }
    84  
    85  func (d *Dumper) Close() {
    86  	close(d.stop)
    87  	d.wg.Wait()
    88  }
    89  
    90  func (d *Dumper) doDumpFile() {
    91  	var err error
    92  	fileName := d.prefix + ".dump"
    93  	logger.Info(context.Background(), "http server start dump data", fileName)
    94  	var f *os.File
    95  	closeFile := func() {
    96  		if f != nil {
    97  			_ = f.Close()
    98  		}
    99  	}
   100  	cutFile := func() (f *os.File, err error) {
   101  		nFile := path.Join(path.Join(config.LoongcollectorGlobalConfig.LoongCollectorDebugDir, "dump"), fileName+"_"+time.Now().Format("2006-01-02_15"))
   102  		if len(d.dumpDataKeepFiles) == 0 || d.dumpDataKeepFiles[len(d.dumpDataKeepFiles)-1] != nFile {
   103  			d.dumpDataKeepFiles = append(d.dumpDataKeepFiles, nFile)
   104  		}
   105  		if len(d.dumpDataKeepFiles) > d.maxKeepFiles {
   106  			rmLen := len(d.dumpDataKeepFiles) - d.maxKeepFiles
   107  			for i := 0; i < rmLen; i++ {
   108  				_ = os.Remove(d.dumpDataKeepFiles[i])
   109  			}
   110  			d.dumpDataKeepFiles = d.dumpDataKeepFiles[rmLen:]
   111  		}
   112  		closeFile()
   113  		return os.OpenFile(d.dumpDataKeepFiles[len(d.dumpDataKeepFiles)-1], os.O_CREATE|os.O_WRONLY, 0600)
   114  	}
   115  	lastHour := 0
   116  	offset := int64(0)
   117  	for {
   118  		select {
   119  		case data := <-d.input:
   120  			if time.Now().Hour() != lastHour {
   121  				file, cerr := cutFile()
   122  				if cerr != nil {
   123  					logger.Error(context.Background(), "DUMP_FILE_ALARM", "cut new file error", err)
   124  				} else {
   125  					offset, _ = file.Seek(0, io.SeekEnd)
   126  					f = file
   127  				}
   128  			}
   129  			if f != nil {
   130  				d.AddDelta(1)
   131  				buffer := bytes.NewBuffer([]byte{})
   132  				b, _ := json.Marshal(data)
   133  				if err = binary.Write(buffer, binary.BigEndian, uint32(len(b))); err != nil {
   134  					continue
   135  				}
   136  				if _, err = f.WriteAt(buffer.Bytes(), offset); err != nil {
   137  					continue
   138  				}
   139  				if _, err = f.WriteAt(b, offset+4); err != nil {
   140  					continue
   141  				}
   142  				offset += int64(4 + len(b))
   143  			}
   144  		case <-d.stop:
   145  			return
   146  		}
   147  	}
   148  }
   149  
   150  func (d *Dumper) InputChannel() chan *DumpData {
   151  	return d.input
   152  }
   153  
   154  func NewDumper(prefix string, maxFiles int) *Dumper {
   155  	return &Dumper{
   156  		prefix:       prefix,
   157  		maxKeepFiles: maxFiles,
   158  	}
   159  }
   160  
   161  func (d *Dumper) Begin(callback func()) bool {
   162  	d.writeCounter = new(async.UnitTestCounterHelper)
   163  	d.writeCounter.Begin(callback)
   164  	return true
   165  }
   166  
   167  func (d *Dumper) End(timeout time.Duration, expectNum int64) error {
   168  	return d.writeCounter.End(timeout, expectNum)
   169  }
   170  
   171  func (d *Dumper) AddDelta(num int64) {
   172  	if d.writeCounter != nil {
   173  		d.writeCounter.AddDelta(num)
   174  	}
   175  }