| 38 | } |
| 39 | |
| 40 | func Encode(partitionID int32, tenantID string, req *tempopb.PushBytesRequest, maxSize int) ([]*kgo.Record, error) { |
| 41 | reqSize := req.Size() |
| 42 | |
| 43 | // Fast path for small requests |
| 44 | if reqSize <= maxSize { |
| 45 | rec, err := marshalWriteRequestToRecord(partitionID, tenantID, req) |
| 46 | if err != nil { |
| 47 | return nil, err |
| 48 | } |
| 49 | return []*kgo.Record{rec}, nil |
| 50 | } |
| 51 | |
| 52 | var records []*kgo.Record |
| 53 | batch := encoderPool.Get().(*tempopb.PushBytesRequest) |
| 54 | defer encoderPoolPut(batch) |
| 55 | |
| 56 | currentSize := 0 |
| 57 | |
| 58 | for i, entry := range req.Traces { |
| 59 | l := entry.Size() + len(req.Ids[i]) |
| 60 | // Size of the entry in the req |
| 61 | entrySize := 1 + l + sovPush(uint64(l)) |
| 62 | |
| 63 | // Check if a single entry is too big |
| 64 | if entrySize > maxSize || (i == 0 && currentSize+entrySize > maxSize) { |
| 65 | return nil, fmt.Errorf("single entry size (%d) exceeds maximum allowed size (%d)", entrySize, maxSize) |
| 66 | } |
| 67 | |
| 68 | if currentSize+entrySize > maxSize { |
| 69 | // Current req is full, create a record and start a new req |
| 70 | if len(batch.Traces) > 0 { |
| 71 | rec, err := marshalWriteRequestToRecord(partitionID, tenantID, batch) |
| 72 | if err != nil { |
| 73 | return nil, err |
| 74 | } |
| 75 | records = append(records, rec) |
| 76 | } |
| 77 | // Reset currentStream |
| 78 | batch.Traces = batch.Traces[:0] |
| 79 | batch.Ids = batch.Ids[:0] |
| 80 | currentSize = 0 |
| 81 | } |
| 82 | batch.Traces = append(batch.Traces, entry) |
| 83 | batch.Ids = append(batch.Ids, req.Ids[i]) |
| 84 | currentSize += entrySize |
| 85 | } |
| 86 | |
| 87 | // Handle any remaining entries |
| 88 | if len(batch.Traces) > 0 { |
| 89 | rec, err := marshalWriteRequestToRecord(partitionID, tenantID, batch) |
| 90 | if err != nil { |
| 91 | return nil, err |
| 92 | } |
| 93 | records = append(records, rec) |
| 94 | } |
| 95 | |
| 96 | if len(records) == 0 { |
| 97 | return nil, errors.New("no valid records created") |