1/*-
2 * See the file LICENSE for redistribution information.
3 *
4 * Copyright (c) 1997-2009 Oracle.  All rights reserved.
5 *
6 * $Id$
7 */
8
9#include <sys/types.h>
10
11#include <stdio.h>
12#include <stdlib.h>
13#include <string.h>
14
15#ifdef _WIN32
16extern int getopt(int, char * const *, const char *);
17#else
18#include <unistd.h>
19#endif
20
21#include <db.h>
22
23#define	DATABASE	"stream.db"
24#define	CHUNK_SIZE	500
25#define	DATA_SIZE	CHUNK_SIZE * 100
26
27int main __P((int, char *[]));
28int usage __P((void));
29int invarg __P((const char *, int, const char *));
30
31int
32main(argc, argv)
33	int argc;
34	char *argv[];
35{
36	extern char *optarg;
37	extern int optind;
38	DB *dbp;
39	DBC *dbcp;
40	DBT key, data;
41	DBTYPE db_type;
42	int ch, chunk_sz, chunk_off, data_sz, i, ret, rflag;
43	int page_sz;
44	char *database, *buf;
45	const char *progname = "ex_stream";		/* Program name. */
46
47	chunk_sz = CHUNK_SIZE;
48	data_sz = DATA_SIZE;
49	chunk_off = page_sz = rflag = 0;
50	db_type = DB_BTREE;
51	while ((ch = getopt(argc, argv, "c:d:p:t:")) != EOF)
52		switch (ch) {
53		case 'c':
54			if ((chunk_sz = atoi(optarg)) <= 0)
55				return (invarg(progname, ch, optarg));
56			break;
57		case 'd':
58			if ((data_sz = atoi(optarg)) <= 0)
59				return (invarg(progname, ch, optarg));
60			break;
61		case 'p':
62			if ((page_sz = atoi(optarg)) <= 0 ||
63			    page_sz % 2 != 0 || page_sz < 512 ||
64			    page_sz > 64 * 1024)
65				return (invarg(progname, ch, optarg));
66			break;
67		case 't':
68			switch (optarg[0]) {
69			case 'b':
70				db_type = DB_BTREE;
71				break;
72			case 'h':
73				db_type = DB_HASH;
74				break;
75			case 'r':
76				db_type = DB_RECNO;
77				break;
78			default:
79				return (invarg(progname, ch, optarg));
80				break;
81			}
82			break;
83		case '?':
84		default:
85			return (usage());
86		}
87	argc -= optind;
88	argv += optind;
89
90	/* Accept optional database name. */
91	database = *argv == NULL ? DATABASE : argv[0];
92
93	if (chunk_sz > data_sz) {
94		fprintf(stderr,
95"Chunk size must be less than and a factor of the data size\n");
96
97		return (usage());
98	}
99
100	/* Discard any existing database. */
101	(void)remove(database);
102
103	/* Create and initialize database object, open the database. */
104	if ((ret = db_create(&dbp, NULL, 0)) != 0) {
105		fprintf(stderr,
106		    "%s: db_create: %s\n", progname, db_strerror(ret));
107		return (EXIT_FAILURE);
108	}
109	dbp->set_errfile(dbp, stderr);
110	dbp->set_errpfx(dbp, progname);
111	if (page_sz != 0 && (ret = dbp->set_pagesize(dbp, page_sz)) != 0) {
112		dbp->err(dbp, ret, "set_pagesize");
113		goto err1;
114	}
115	if ((ret = dbp->set_cachesize(dbp, 0, 32 * 1024, 0)) != 0) {
116		dbp->err(dbp, ret, "set_cachesize");
117		goto err1;
118	}
119	if ((ret = dbp->open(dbp,
120	    NULL, database, NULL, db_type, DB_CREATE, 0664)) != 0) {
121		dbp->err(dbp, ret, "%s: open", database);
122		goto err1;
123	}
124
125	/* Ensure the data size is a multiple of the chunk size. */
126	data_sz = data_sz - (data_sz % chunk_sz);
127
128	/* Initialize the key/data pair for a streaming insert. */
129	memset(&key, 0, sizeof(key));
130	memset(&data, 0, sizeof(data));
131	key.data = &chunk_sz; /* Our key value does not really matter. */
132	key.size = sizeof(int);
133	data.ulen = data_sz;
134	data.size = chunk_sz;
135	data.data = buf = malloc(data_sz);
136	data.flags = DB_DBT_USERMEM | DB_DBT_PARTIAL;
137
138	/* Populate the data with something. */
139	for (i = 0; i < data_sz; ++i)
140		buf[i] = (char)('a' + i % ('z' - 'a'));
141
142	if ((ret = dbp->cursor(dbp, NULL, &dbcp, 0)) != 0) {
143		dbp->err(dbp, ret, "DB->cursor");
144		goto err1;
145	}
146	for (chunk_off = 0; chunk_off < data_sz; chunk_off += chunk_sz) {
147		data.size = chunk_sz;
148		if ((ret = dbcp->put(dbcp, &key, &data,
149		    (chunk_off == 0 ? DB_KEYFIRST : DB_CURRENT)) != 0)) {
150			dbp->err(dbp, ret, "DBCursor->put");
151			goto err2;
152		}
153		data.doff += chunk_sz;
154	}
155	if ((ret = dbcp->close(dbcp)) != 0) {
156		dbp->err(dbp, ret, "DBcursor->close");
157		goto err1;
158	}
159
160	memset(data.data, 0, data.ulen);
161	/* Retrieve the data item in chunks. */
162	if ((ret = dbp->cursor(dbp, NULL, &dbcp, 0)) != 0) {
163		dbp->err(dbp, ret, "DB->cursor");
164		goto err1;
165	}
166	data.doff = 0;
167	data.dlen = chunk_sz;
168	memset(data.data, 0, data.ulen);
169
170	/*
171	 * Loop over the item, retrieving a chunk at a time.
172	 * The requested chunk will be stored at the start of data.data.
173	 */
174	for (chunk_off = 0; chunk_off < data_sz; chunk_off += chunk_sz) {
175		if ((ret = dbcp->get(dbcp, &key, &data,
176		    (chunk_off == 0 ? DB_SET : DB_CURRENT)) != 0)) {
177			dbp->err(dbp, ret, "DBCursor->get");
178			goto err2;
179		}
180		data.doff += chunk_sz;
181	}
182
183	if ((ret = dbcp->close(dbcp)) != 0) {
184		dbp->err(dbp, ret, "DBcursor->close");
185		goto err1;
186	}
187	if ((ret = dbp->close(dbp, 0)) != 0) {
188		fprintf(stderr,
189		    "%s: DB->close: %s\n", progname, db_strerror(ret));
190		return (EXIT_FAILURE);
191	}
192	return (EXIT_SUCCESS);
193
194err2:	(void)dbcp->close(dbcp);
195err1:	(void)dbp->close(dbp, 0);
196	return (EXIT_FAILURE);
197}
198
199int
200invarg(progname, arg, str)
201	const char *progname;
202	int arg;
203	const char *str;
204{
205	(void)fprintf(stderr,
206	    "%s: invalid argument for -%c: %s\n", progname, arg, str);
207	return (EXIT_FAILURE);
208}
209
210int
211usage()
212{
213	(void)fprintf(stderr,
214"usage: ex_stream [-c int] [-d int] [-p int] [-t char] [database]\n");
215	(void)fprintf(stderr, "Where options are:\n");
216	(void)fprintf(stderr, "\t-c set the chunk size.\n");
217	(void)fprintf(stderr, "\t-d set the total record size.\n");
218	(void)fprintf(stderr,
219	    "\t-t choose a database type btree (b), hash (h) or recno (r)\n");
220
221	return (EXIT_FAILURE);
222}
223