github.com/lirm/aeron-go@v0.0.0-20230415210743-920325491dc4/cluster/snapshot_loader.go (about)

     1  // Licensed under the Apache License, Version 2.0 (the "License");
     2  // you may not use this file except in compliance with the License.
     3  // You may obtain a copy of the License at
     4  //
     5  // http://www.apache.org/licenses/LICENSE-2.0
     6  //
     7  // Unless required by applicable law or agreed to in writing, software
     8  // distributed under the License is distributed on an "AS IS" BASIS,
     9  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    10  // See the License for the specific language governing permissions and
    11  // limitations under the License.
    12  
    13  package cluster
    14  
    15  import (
    16  	"bytes"
    17  	"fmt"
    18  
    19  	"github.com/lirm/aeron-go/aeron"
    20  	"github.com/lirm/aeron-go/aeron/atomic"
    21  	"github.com/lirm/aeron-go/aeron/logbuffer"
    22  	"github.com/lirm/aeron-go/cluster/codecs"
    23  )
    24  
    25  type snapshotLoader struct {
    26  	agent      *ClusteredServiceAgent
    27  	img        aeron.Image
    28  	marshaller *codecs.SbeGoMarshaller
    29  	isDone     bool
    30  	inSnapshot bool
    31  	appVersion int32
    32  	timeUnit   codecs.ClusterTimeUnitEnum
    33  	err        error
    34  }
    35  
    36  func newSnapshotLoader(agent *ClusteredServiceAgent, img aeron.Image) *snapshotLoader {
    37  	return &snapshotLoader{agent: agent, img: img, marshaller: codecs.NewSbeGoMarshaller()}
    38  }
    39  
    40  func (loader *snapshotLoader) poll() int {
    41  	if loader.isDone {
    42  		return 0
    43  	}
    44  	return loader.img.Poll(loader.onFragment, 1)
    45  }
    46  
    47  func (loader *snapshotLoader) onFragment(
    48  	buffer *atomic.Buffer,
    49  	offset int32,
    50  	length int32,
    51  	header *logbuffer.Header,
    52  ) {
    53  	if length < SBEHeaderLength {
    54  		return
    55  	}
    56  	blockLength := buffer.GetUInt16(offset)
    57  	templateId := buffer.GetUInt16(offset + 2)
    58  	schemaId := buffer.GetUInt16(offset + 4)
    59  	version := buffer.GetUInt16(offset + 6)
    60  	if schemaId != ClusterSchemaId {
    61  		logger.Errorf("SnapshotLoader: unexpected schemaId=%d templateId=%d blockLen=%d version=%d",
    62  			schemaId, templateId, blockLength, version)
    63  		return
    64  	}
    65  	offset += SBEHeaderLength
    66  	length -= SBEHeaderLength
    67  
    68  	buf := &bytes.Buffer{}
    69  	buffer.WriteBytes(buf, offset, length)
    70  
    71  	switch templateId {
    72  	case snapshotMarkerTemplateId:
    73  		marker := &codecs.SnapshotMarker{}
    74  		if err := marker.Decode(loader.marshaller, buf, version, blockLength, true); err != nil {
    75  			loader.err = err
    76  			loader.isDone = true
    77  		} else if marker.TypeId != snapshotTypeId {
    78  			loader.err = fmt.Errorf("unexpected snapshot type: %d", marker.TypeId)
    79  		} else if marker.Mark == codecs.SnapshotMark.BEGIN {
    80  			if loader.inSnapshot {
    81  				loader.err = fmt.Errorf("already in snapshot, pos=%d", header.Position())
    82  				loader.isDone = true
    83  			} else {
    84  				loader.inSnapshot = true
    85  				loader.appVersion = marker.AppVersion
    86  				loader.timeUnit = marker.TimeUnit
    87  			}
    88  		} else if marker.Mark == codecs.SnapshotMark.END {
    89  			if !loader.inSnapshot {
    90  				loader.err = fmt.Errorf("missing beging snapshot, pos=%d", header.Position())
    91  				loader.isDone = true
    92  			} else {
    93  				loader.inSnapshot = false
    94  				loader.isDone = true
    95  			}
    96  		} else {
    97  			loader.err = fmt.Errorf("unexpected snapshot mark, pos=%d inSnapshot=%v mark=%v", header.Position(), loader.inSnapshot, marker.Mark)
    98  			loader.isDone = true
    99  		}
   100  	case clientSessionTemplateId:
   101  		s := codecs.ClientSession{}
   102  		if err := s.Decode(loader.marshaller, buf, version, blockLength, true); err != nil {
   103  			loader.err = err
   104  			loader.isDone = true
   105  		} else {
   106  			container, err := newContainerClientSession(
   107  				s.ClusterSessionId, s.ResponseStreamId, string(s.ResponseChannel), s.EncodedPrincipal, loader.agent)
   108  			if err == nil {
   109  				loader.agent.addSessionFromSnapshot(container)
   110  			} else {
   111  				loader.err = err
   112  				loader.isDone = true
   113  			}
   114  		}
   115  	default:
   116  		logger.Debugf("SnapshotLoader: unknown templateId=%d at pos=%d", templateId, header.Position())
   117  	}
   118  }