go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/logdog/client/butler/output/directory/stream.go (about)

     1  // Copyright 2019 The LUCI 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 directory
    16  
    17  import (
    18  	"fmt"
    19  	"os"
    20  	"path/filepath"
    21  
    22  	"github.com/golang/protobuf/jsonpb"
    23  
    24  	"go.chromium.org/luci/common/errors"
    25  	"go.chromium.org/luci/logdog/api/logpb"
    26  )
    27  
    28  // stream is the stateful output for a single log stream.
    29  type stream struct {
    30  	curFile *os.File // nil if no file open
    31  
    32  	basePath      string
    33  	fname         string
    34  	isDatagram    bool
    35  	datagramCount int
    36  }
    37  
    38  func newStream(basePath string, desc *logpb.LogStreamDescriptor) (*stream, error) {
    39  	relPath := filepath.Clean(desc.Name)
    40  	dir, fname := filepath.Split(relPath)
    41  	basePath = filepath.Join(basePath, dir)
    42  
    43  	_ = os.MkdirAll(basePath, 0750)
    44  	metaF, err := os.Create(filepath.Join(basePath, ".meta."+fname))
    45  	if err != nil {
    46  		return nil, errors.Annotate(err, "opening meta file for %s", relPath).Err()
    47  	}
    48  	defer metaF.Close()
    49  	err = (&jsonpb.Marshaler{
    50  		Indent:   "  ",
    51  		OrigName: true,
    52  	}).Marshal(metaF, desc)
    53  	if err != nil {
    54  		return nil, errors.Annotate(err, "writing meta file for %s", relPath).Err()
    55  	}
    56  
    57  	ret := stream{basePath: basePath, fname: fname}
    58  	if desc.StreamType == logpb.StreamType_DATAGRAM {
    59  		ret.isDatagram = true
    60  	} else {
    61  		ret.curFile, err = os.Create(filepath.Join(basePath, fname))
    62  	}
    63  	return &ret, err
    64  }
    65  
    66  func (s *stream) getCurFile() (*os.File, error) {
    67  	if s.curFile != nil {
    68  		return s.curFile, nil
    69  	}
    70  
    71  	if !s.isDatagram {
    72  		return nil, errors.New(
    73  			"cannot call getCurFile for a non-datagram with a closed file")
    74  	}
    75  
    76  	var err error
    77  	s.curFile, err = os.Create(
    78  		filepath.Join(s.basePath, fmt.Sprintf("_%05d.%s", s.datagramCount, s.fname)))
    79  	if err != nil {
    80  		return nil, errors.Annotate(err, "could not open %d'th datagram of %s",
    81  			s.datagramCount, filepath.Join(s.basePath, s.fname)).Err()
    82  	}
    83  	s.datagramCount++
    84  
    85  	return s.curFile, nil
    86  }
    87  
    88  func (s *stream) closeCurFile() {
    89  	if s.curFile != nil {
    90  		s.curFile.Close()
    91  		s.curFile = nil
    92  	}
    93  }
    94  
    95  // ingestBundleEntry writes the data from `be` to disk
    96  //
    97  // Returns closed == true if `be` was terminal and the stream can be closed now.
    98  func (s *stream) ingestBundleEntry(be *logpb.ButlerLogBundle_Entry) (closed bool, err error) {
    99  	for _, le := range be.GetLogs() {
   100  		curFile, err := s.getCurFile()
   101  		if err != nil {
   102  			return false, err
   103  		}
   104  
   105  		switch x := le.Content.(type) {
   106  		case *logpb.LogEntry_Datagram:
   107  			dg := x.Datagram
   108  			_, err = s.curFile.Write(dg.Data)
   109  			if err == nil {
   110  				if dg.Partial == nil || dg.Partial.Last {
   111  					s.closeCurFile()
   112  				}
   113  			}
   114  		case *logpb.LogEntry_Text:
   115  			for _, line := range x.Text.Lines {
   116  				_, err = curFile.Write(line.Value)
   117  				if err == nil {
   118  					_, err = curFile.WriteString("\n")
   119  				}
   120  			}
   121  		case *logpb.LogEntry_Binary:
   122  			_, err = curFile.Write(x.Binary.Data)
   123  		}
   124  
   125  		if err != nil {
   126  			return false, err
   127  		}
   128  	}
   129  	if be.Terminal {
   130  		s.Close()
   131  		return true, nil
   132  	}
   133  	return false, nil
   134  }
   135  
   136  func (s *stream) Close() {
   137  	s.closeCurFile()
   138  }