| 1 | package agent |
| 2 | |
| 3 | import ( |
| 4 | "context" |
| 5 | "encoding/json" |
| 6 | "fmt" |
| 7 | "os" |
| 8 | "path/filepath" |
| 9 | |
| 10 | "reasonix/internal/fileutil" |
| 11 | "reasonix/internal/provider" |
| 12 | ) |
| 13 | |
| 14 | const sessionContextWriteChunk = 256 << 10 |
| 15 | |
| 16 | type jsonMarshalResult struct { |
| 17 | data []byte |
| 18 | err error |
| 19 | } |
| 20 | |
| 21 | type sessionPublishStartHookKey struct{} |
| 22 | |
| 23 | func withSessionPublishStartHook(ctx context.Context, hook func()) context.Context { |
| 24 | return context.WithValue(ctx, sessionPublishStartHookKey{}, hook) |
| 25 | } |
| 26 | |
| 27 | // marshalJSONContext lets maintenance work release session locks promptly |
| 28 | // when a large attachment is still being encoded. Foreground saves keep the |
| 29 | // synchronous path through their background context wrappers. |
| 30 | func marshalJSONContext(ctx context.Context, value any) ([]byte, error) { |
| 31 | return marshalJSONWithIndentContext(ctx, value, false) |
| 32 | } |
| 33 | |
| 34 | func marshalJSONIndentContext(ctx context.Context, value any) ([]byte, error) { |
| 35 | return marshalJSONWithIndentContext(ctx, value, true) |
| 36 | } |
| 37 | |
| 38 | func marshalJSONWithIndentContext(ctx context.Context, value any, indent bool) ([]byte, error) { |
| 39 | if err := ctx.Err(); err != nil { |
| 40 | return nil, err |
| 41 | } |
| 42 | marshal := func() ([]byte, error) { |
| 43 | if indent { |
| 44 | return json.MarshalIndent(value, "", " ") |
| 45 | } |
| 46 | return json.Marshal(value) |
| 47 | } |
| 48 | if ctx.Done() == nil { |
| 49 | return marshal() |
| 50 | } |
| 51 | resultCh := make(chan jsonMarshalResult, 1) |
| 52 | go func() { |
| 53 | data, err := marshal() |
| 54 | resultCh <- jsonMarshalResult{data: data, err: err} |
| 55 | }() |
| 56 | select { |
| 57 | case <-ctx.Done(): |
| 58 | return nil, ctx.Err() |
| 59 | case result := <-resultCh: |
| 60 | return result.data, result.err |
| 61 | } |
| 62 | } |
| 63 | |
| 64 | func atomicWriteFileContext(ctx context.Context, path, pattern, crashOp string, data []byte, perm os.FileMode, syncFile bool) error { |
| 65 | if err := ctx.Err(); err != nil { |
| 66 | return err |
| 67 | } |
| 68 | if crashOp != "" { |
| 69 | fileutil.Crash(crashOp, path) |
| 70 | } |
| 71 | if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { |
| 72 | return err |
| 73 | } |
| 74 | tmp, err := os.CreateTemp(filepath.Dir(path), pattern) |
| 75 | if err != nil { |
| 76 | return err |
| 77 | } |
| 78 | tmpPath := tmp.Name() |
| 79 | cleanup := func() { |
| 80 | _ = tmp.Close() |
| 81 | _ = os.Remove(tmpPath) |
| 82 | } |
| 83 | if err := tmp.Chmod(perm); err != nil { |
| 84 | cleanup() |
| 85 | return err |
| 86 | } |
| 87 | if err := writeContextBytes(ctx, tmp, path, data); err != nil { |
| 88 | cleanup() |
| 89 | return err |
| 90 | } |
| 91 | if err := ctx.Err(); err != nil { |
| 92 | cleanup() |
| 93 | return err |
| 94 | } |
| 95 | if syncFile { |
| 96 | if err := tmp.Sync(); err != nil { |
| 97 | cleanup() |
| 98 | return err |
| 99 | } |
| 100 | } |
| 101 | if err := tmp.Close(); err != nil { |
| 102 | _ = os.Remove(tmpPath) |
| 103 | return err |
| 104 | } |
| 105 | if err := ctx.Err(); err != nil { |
| 106 | _ = os.Remove(tmpPath) |
| 107 | return err |
| 108 | } |
| 109 | if err := fileutil.ReplaceFile(tmpPath, path); err != nil { |
| 110 | _ = os.Remove(tmpPath) |
| 111 | return err |
| 112 | } |
| 113 | return nil |
| 114 | } |
| 115 | |
| 116 | func writeContextBytes(ctx context.Context, file *os.File, path string, data []byte) error { |
| 117 | for len(data) > 0 { |
| 118 | if err := ctx.Err(); err != nil { |
| 119 | return err |
| 120 | } |
| 121 | chunk := min(len(data), sessionContextWriteChunk) |
| 122 | written, writeErr := file.Write(data[:chunk]) |
| 123 | if writeErr != nil { |
| 124 | return writeErr |
| 125 | } |
| 126 | if written == 0 { |
| 127 | return fmt.Errorf("write %s: no progress", path) |
| 128 | } |
| 129 | data = data[written:] |
| 130 | } |
| 131 | return ctx.Err() |
| 132 | } |
| 133 | |
| 134 | func writeSessionMessages(path string, msgs []provider.Message) error { |
| 135 | return writeSessionMessagesContext(context.Background(), path, msgs) |
| 136 | } |
| 137 | |
| 138 | func writeSessionMessagesContext(ctx context.Context, path string, msgs []provider.Message) error { |
| 139 | // The compatibility transcript is a crash-safe anchor when the event log |
| 140 | // is damaged. Maintenance checks cancellation before every publish step. |
| 141 | if err := ctx.Err(); err != nil { |
| 142 | return err |
| 143 | } |
| 144 | fileutil.Crash("session-checkpoint", path) |
| 145 | if err := os.MkdirAll(filepath.Dir(path), 0o755); err != nil { |
| 146 | return err |
| 147 | } |
| 148 | tmp, err := os.CreateTemp(filepath.Dir(path), ".session.*.tmp") |
| 149 | if err != nil { |
| 150 | return fmt.Errorf("create session tmp: %w", err) |
| 151 | } |
| 152 | tmpPath := tmp.Name() |
| 153 | cleanup := func() { |
| 154 | _ = tmp.Close() |
| 155 | _ = os.Remove(tmpPath) |
| 156 | } |
| 157 | if hook, _ := ctx.Value(sessionPublishStartHookKey{}).(func()); hook != nil { |
| 158 | hook() |
| 159 | } |
| 160 | for _, message := range msgs { |
| 161 | data, err := marshalJSONContext(ctx, message) |
| 162 | if err != nil { |
| 163 | cleanup() |
| 164 | return fmt.Errorf("encode message: %w", err) |
| 165 | } |
| 166 | data = append(data, '\n') |
| 167 | if err := writeContextBytes(ctx, tmp, path, data); err != nil { |
| 168 | cleanup() |
| 169 | return fmt.Errorf("write session messages: %w", err) |
| 170 | } |
| 171 | } |
| 172 | if err := ctx.Err(); err != nil { |
| 173 | cleanup() |
| 174 | return err |
| 175 | } |
| 176 | if err := tmp.Sync(); err != nil { |
| 177 | cleanup() |
| 178 | return err |
| 179 | } |
| 180 | if err := tmp.Close(); err != nil { |
| 181 | _ = os.Remove(tmpPath) |
| 182 | return err |
| 183 | } |
| 184 | if err := ctx.Err(); err != nil { |
| 185 | _ = os.Remove(tmpPath) |
| 186 | return err |
| 187 | } |
| 188 | if err := fileutil.ReplaceFile(tmpPath, path); err != nil { |
| 189 | _ = os.Remove(tmpPath) |
| 190 | return fmt.Errorf("write session messages: %w", err) |
| 191 | } |
| 192 | return nil |
| 193 | } |
| 194 |